Add noise_loss param for deprecated batch loaders

EmmaRenauld · Sep 26, 2023 · a3a2c32 · a3a2c32
1 parent 8cf0e82
commit a3a2c32
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/dwi_ml/training/batch_loaders.py b/dwi_ml/training/batch_loaders.py
@@ -161,6 +161,12 @@ def params_for_checkpoint(self):
     @classmethod
     def init_from_checkpoint(cls, dataset, model, checkpoint_state,
                              new_log_level):
+        # Adding noise_gaussian_size_loss for deprecated batch loaders
+        if 'noise_gaussian_size_loss' not in checkpoint_state:
+            logging.warning("Deprecated batch loader. Did not contain a "
+                            "noise_gaussian_size_loss value. Setting to 0.0.")
+            checkpoint_state['noise_gaussian_size_loss'] = 0.0
+
         batch_loader = cls(dataset=dataset, model=model,
                            log_level=new_log_level, **checkpoint_state)
         return batch_loader