Merge pull request #107 from HannesStark/training-fixes

training fixes

src/boltzgen/model/validation/refolding.py

@@ -111,6 +111,7 @@ class RefoldingValidator(design.DesignValidator):
                 strict=False,
                 map_location=model.device,
                 predict_args=self.folding_args,
+                weights_only=False,
                 **self.folding_model_args,
             )
             model_module.eval()

src/boltzgen/resources/config/train/boltzgen.yaml

@@ -428,7 +428,6 @@ model:
           num_workers: 4
           pin_memory: true
         return_native: true
-      predict_task: null
 
     folding_checkpoint: ./training_data/boltz2_fold.ckpt
 

src/boltzgen/task/train/train.py

@@ -140,7 +140,7 @@ class Training(Task):
 
             print(f"Loading model from {file_path}")
             model_module = type(model_module).load_from_checkpoint(
-                file_path, map_location="cpu", strict=False, **(model_module.hparams)
+                file_path, map_location="cpu", strict=False, weights_only=False, **(model_module.hparams)
             )
 
         # Create checkpoint callback