Refactor input_channels assignment in SegmentationDOFA.export_model()

NRCan · Oct 24, 2024 · 5f3943c · 5f3943c
1 parent cb5de25
commit 5f3943c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/geo_deep_learning/tasks_with_models/segmentation_dofa.py b/geo_deep_learning/tasks_with_models/segmentation_dofa.py
@@ -80,7 +80,7 @@ def on_train_end(self):
                 self.export_model(best_model_path, best_model_export_path, self.trainer.datamodule)
 
     def export_model(self, checkpoint_path: str, export_path: str, datamodule: LightningDataModule):
-        input_channels = self.hparams["init_args"]["in_channels"]
+        input_channels = self.hparams["in_channels"]
         map_location = "cuda"
         if self.device.type == "cpu":
             map_location = "cpu"