add predict_step to MultiLabelClassificationTask

tests/trainers/test_classification.py

@@ -153,6 +153,7 @@ def test_trainer(
         trainer = Trainer(fast_dev_run=True, log_every_n_steps=1, max_epochs=1)
         trainer.fit(model=model, datamodule=datamodule)
         trainer.test(model=model, datamodule=datamodule)
+        trainer.predict(model=model, dataloaders=datamodule.val_dataloader())
 
     def test_no_logger(self) -> None:
         conf = OmegaConf.load(os.path.join("tests", "conf", "bigearthnet_s1.yaml"))

torchgeo/trainers/classification.py

@@ -388,3 +388,15 @@ def test_step(self, *args: Any, **kwargs: Any) -> None:
         # by default, the test and validation steps only log per *epoch*
         self.log("test_loss", loss, on_step=False, on_epoch=True)
         self.test_metrics(y_hat_hard, y)
+
+    def predict_step(self, *args: Any, **kwargs: Any) -> Tensor:
+        """Compute and return the predictions.
+        Args:
+            batch: the output of your DataLoader
+        Returns:
+            predicted sigmoid probabilities
+        """
+        batch = args[0]
+        x = batch["image"]
+        y_hat = torch.sigmoid(self(x))
+        return y_hat