fix: type issues

luxonis · Oct 14, 2024 · 01918f0 · 01918f0
1 parent 22bba41
commit 01918f0
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/luxonis_train/callbacks/ema.py b/luxonis_train/callbacks/ema.py
@@ -171,7 +171,8 @@ def on_validation_epoch_start(
 
         self.collected_state_dict = deepcopy(pl_module.state_dict())
 
-        pl_module.load_state_dict(self.ema.state_dict_ema)
+        if self.ema is not None:
+            pl_module.load_state_dict(self.ema.state_dict_ema)
 
     def on_validation_end(
         self, trainer: pl.Trainer, pl_module: pl.LightningModule
@@ -183,7 +184,8 @@ def on_validation_end(
         @type pl_module: L{pl.LightningModule}
         @param pl_module: Pytorch Lightning module.
         """
-        pl_module.load_state_dict(self.collected_state_dict)
+        if self.collected_state_dict is not None:
+            pl_module.load_state_dict(self.collected_state_dict)
 
     def on_train_end(
         self, trainer: pl.Trainer, pl_module: pl.LightningModule
@@ -195,8 +197,9 @@ def on_train_end(
         @type pl_module: L{pl.LightningModule}
         @param pl_module: Pytorch Lightning module.
         """
-        pl_module.load_state_dict(self.ema.state_dict_ema)
-        logger.info("Model weights replaced with the EMA weights.")
+        if self.ema is not None:
+            pl_module.load_state_dict(self.ema.state_dict_ema)
+            logger.info("Model weights replaced with the EMA weights.")
 
     def on_save_checkpoint(
         self,
@@ -213,12 +216,14 @@ def on_save_checkpoint(
         @type checkpoint: dict
         @param checkpoint: Pytorch Lightning checkpoint.
         """
-        checkpoint["state_dict"] = self.ema.state_dict_ema
+        if self.ema is not None:
+            checkpoint["state_dict"] = self.ema.state_dict_ema
 
     def on_load_checkpoint(self, callback_state: dict) -> None:
         """Load the EMA state_dict from the checkpoint.
 
         @type callback_state: dict
         @param callback_state: Pytorch Lightning callback state.
         """
-        self.ema.state_dict_ema = callback_state["state_dict"]
+        if self.ema is not None:
+            self.ema.state_dict_ema = callback_state["state_dict"]