Fix best models loading (for dataset generation)

PySymGym · Dec 25, 2023 · 9e32202 · 9e32202
1 parent 136bc2d
commit 9e32202
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/VSharp.ML.AIAgent/ml/common_model/utils.py b/VSharp.ML.AIAgent/ml/common_model/utils.py
@@ -63,12 +63,12 @@ def csv2best_models(ref_model_init: t.Callable[[], torch.nn.Module]):
                 )
                 ref_model = load_model(Path(path_to_model), model=ref_model_init())
 
-            ref_model.to(GeneralConfig.DEVICE)
-            best_models[map_name] = (
-                ref_model,
-                best_model_score[map_name],
-                best_model_name,
-            )
+                ref_model.to(GeneralConfig.DEVICE)
+                best_models[map_name] = (
+                    ref_model,
+                    best_model_score[map_name],
+                    best_model_name,
+                )
     return best_models