Update stopping.py

jwdink · jwdink · commit bcb2738ea436 · 2025-04-21T11:08:50.000-05:00
diff --git a/torchcast/utils/stopping.py b/torchcast/utils/stopping.py
@@ -53,7 +53,7 @@ def from_dict(cls, **kwargs) -> 'Stopping':
         return cls(**kwargs)
 
     @torch.inference_mode()
-    def _get_new_values(self, loss: Optional[float]):
+    def _get_new_values(self, loss: Optional[float]) -> torch.Tensor:
         flat = []
 
         if self.monitor_params:
