fix mypy

microsoft · Apr 13, 2023 · dd2f9a0 · dd2f9a0
1 parent 43c5db8
commit dd2f9a0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchgeo/trainers/regression.py b/torchgeo/trainers/regression.py
@@ -266,6 +266,7 @@ def config_task(self) -> None:
                 f"Currently, only supports 'unet', 'deeplabv3+' and 'fcn'."
             )
 
+        self.loss: nn.Module
         if self.hyperparams["loss"] == "mse":
             self.loss = nn.MSELoss()
         elif self.hyperparams["loss"] == "mae":
@@ -295,7 +296,7 @@ def training_step(self, *args: Any, **kwargs: Any) -> Tensor:
         self.log("train_loss", loss)  # logging to TensorBoard
         self.train_metrics(y_hat, y)
 
-        return loss
+        return cast(Tensor, loss)
 
     def validation_step(self, *args: Any, **kwargs: Any) -> None:
         """Compute validation loss and log example predictions.