Update type hint

Author: gmertes

models/src/anemoi/models/models/encoder_processor_decoder.py

@@ -160,17 +160,17 @@ def _assert_valid_sharding(
 
     def forward(
         self,
-        x: Tensor,
+        x: dict[str, Tensor],
         *,
         model_comm_group: Optional[ProcessGroup] = None,
         grid_shard_shapes: dict[str, list] | None = None,
         **kwargs,
-    ) -> Tensor:
+    ) -> dict[str, Tensor]:
         """Forward pass of the model.
 
         Parameters
         ----------
-        x : Tensor
+        x : dict[str, Tensor]
             Input data
         model_comm_group : Optional[ProcessGroup], optional
             Model communication group, by default None
@@ -179,7 +179,7 @@ def forward(
 
         Returns
         -------
-        Tensor
+        dict[str, Tensor]
             Output of the model, with the same shape as the input (sharded if input is sharded)
         """
         # Multi-dataset case