restrict lt version, update to latest lt version

nndet/ptmodule/base_module.py

@@ -104,11 +104,11 @@ class LightningBaseModule(pl.LightningModule):
         """
         return torch.zeros(*self.example_input_array_shape)
 
-    def summarize(self, mode: Optional[str]) -> Optional[ModelSummary]:
+    def summarize(self, *args, **kwargs) -> Optional[ModelSummary]:
         """
         Save model summary as txt
         """
-        summary = super().summarize(mode=mode)
+        summary = super().summarize(*args, **kwargs)
         save_txt(summary, "./network")
         return summary
 

requirements.txt

-pytorch_lightning
+pytorch_lightning>=1.3.1,<=1.4.2
 scipy
 scikit-learn
 scikit-image>=0.14

scripts/train.py

@@ -283,7 +283,7 @@ def _train(
         weights_summary='full',
         plugins=plugins,
         terminate_on_nan=True,  # TODO: make modular
-        move_metrics_to_cpu=True,
+        move_metrics_to_cpu=False,
         **trainer_kwargs
     )
     trainer.fit(module, datamodule=datamodule)