changed type of key

Author: ghanvert

src/accmt/modules.py

@@ -91,7 +91,7 @@ def training_step(self, batch: Any) -> torch.Tensor:
         """Defines the training logic. Must return a loss tensor (scalar)."""
 
     @override
-    def validation_step(self, key: Any, batch: Any) -> dict:
+    def validation_step(self, key: str, batch: Any) -> dict:
         """
         Defines the validation logic. Must return a dictionary containing
         each metric with predictions and targets, and also the loss value in the dictionary.