update general target script to new logger

spozdn · Jun 21, 2024 · de1d0ab · de1d0ab
1 parent 79d80e8
commit de1d0ab
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/train_model_general_target.py b/src/train_model_general_target.py
@@ -131,7 +131,11 @@ def main():
         load_checkpoint(model, optim, scheduler, f"results/{name_to_load}/checkpoint")
 
     history = []
-    logger = FullLogger(FITTING_SCHEME.SUPPORT_MISSING_VALUES)
+    logger = FullLogger(
+        FITTING_SCHEME.SUPPORT_MISSING_VALUES,
+        FITTING_SCHEME.USE_SHIFT_AGNOSTIC_LOSS,
+        device,
+    )
     mae_model_keeper = ModelKeeper()
     rmse_model_keeper = ModelKeeper()