added **kwargs to predict step

ecmwf · Jan 15, 2025 · fdd92e6 · fdd92e6
1 parent be9460d
commit fdd92e6
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/models/src/anemoi/models/interface/__init__.py b/models/src/anemoi/models/interface/__init__.py
@@ -96,7 +96,9 @@ def _build_model(self) -> None:
         # Use the forward method of the model directly
         self.forward = self.model.forward
 
-    def predict_step(self, batch: torch.Tensor, model_comm_group: Optional[ProcessGroup] = None) -> torch.Tensor:
+    def predict_step(
+        self, batch: torch.Tensor, model_comm_group: Optional[ProcessGroup] = None, **kwargs
+    ) -> torch.Tensor:
         """Prediction step for the model.
 
         Parameters