Fix all Trainers to support PyTorch DataLoaders (#1703)

* Fix all Trainers to support PyTorch DataLoaders * Add DataLoader Fix For MultiLabelClassification --------- Co-authored-by: Kay Benkert <[email protected]>
microsoft · Oct 28, 2023 · 906e113 · 906e113
1 parent 01fe8d5
commit 906e113
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 0 deletions.
diff --git a/torchgeo/trainers/classification.py b/torchgeo/trainers/classification.py
@@ -190,6 +190,7 @@ def validation_step(
         if (
             batch_idx < 10
             and hasattr(self.trainer, "datamodule")
+            and hasattr(self.trainer.datamodule, "plot")
             and self.logger
             and hasattr(self.logger, "experiment")
             and hasattr(self.logger.experiment, "add_figure")
@@ -313,6 +314,7 @@ def validation_step(
         if (
             batch_idx < 10
             and hasattr(self.trainer, "datamodule")
+            and hasattr(self.trainer.datamodule, "plot")
             and self.logger
             and hasattr(self.logger, "experiment")
             and hasattr(self.logger.experiment, "add_figure")

diff --git a/torchgeo/trainers/detection.py b/torchgeo/trainers/detection.py
@@ -259,6 +259,7 @@ def validation_step(
         if (
             batch_idx < 10
             and hasattr(self.trainer, "datamodule")
+            and hasattr(self.trainer.datamodule, "plot")
             and self.logger
             and hasattr(self.logger, "experiment")
             and hasattr(self.logger.experiment, "add_figure")

diff --git a/torchgeo/trainers/regression.py b/torchgeo/trainers/regression.py
@@ -185,6 +185,7 @@ def validation_step(
         if (
             batch_idx < 10
             and hasattr(self.trainer, "datamodule")
+            and hasattr(self.trainer.datamodule, "plot")
             and self.logger
             and hasattr(self.logger, "experiment")
             and hasattr(self.logger.experiment, "add_figure")

diff --git a/torchgeo/trainers/segmentation.py b/torchgeo/trainers/segmentation.py
@@ -245,6 +245,7 @@ def validation_step(
         if (
             batch_idx < 10
             and hasattr(self.trainer, "datamodule")
+            and hasattr(self.trainer.datamodule, "plot")
             and self.logger
             and hasattr(self.logger, "experiment")
             and hasattr(self.logger.experiment, "add_figure")