Fix typing

src/datamodules/mnist_datamodule.py

@@ -59,7 +59,7 @@ def __init__(
         self.data_test: Optional[Dataset] = None
 
     @property
-    def num_classes(self) -> int:
+    def num_classes(self):
         return 10
 
     def prepare_data(self):
@@ -118,11 +118,11 @@ def teardown(self, stage: Optional[str] = None):
         """Clean up after fit or test."""
         pass
 
-    def state_dict(self) -> Dict[str, Any]:
+    def state_dict(self):
         """Extra things to save to checkpoint."""
         return {}
 
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+    def load_state_dict(self, state_dict: Dict[str, Any]):
         """Things to do when loading checkpoint."""
         pass
 

src/tasks/eval_task.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Tuple
+from typing import List, Tuple
 
 import hydra
 from omegaconf import DictConfig
@@ -11,7 +11,7 @@
 
 
 @utils.task_wrapper
-def evaluate(cfg: DictConfig) -> Tuple[None, Dict[str, Any]]:
+def evaluate(cfg: DictConfig) -> Tuple[dict, dict]:
     """Evaluates given checkpoint on a datamodule testset.
 
     This method is wrapped in @task_wrapper decorator which applies extra utilities

src/tasks/train_task.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Tuple
+from typing import List, Tuple
 
 import hydra
 import pytorch_lightning as pl
@@ -12,7 +12,7 @@
 
 
 @utils.task_wrapper
-def train(cfg: DictConfig) -> Tuple[Optional[float], Dict[str, Any]]:
+def train(cfg: DictConfig) -> Tuple[dict, dict]:
     """Trains the model. Can additionally evaluate on a testset, using best weights obtained during
     training.
 

src/utils/utils.py

@@ -84,7 +84,7 @@ def extras(cfg: DictConfig) -> None:
 
 
 @rank_zero_only
-def save_file(path, content) -> None:
+def save_file(path: str, content: str) -> None:
     """Save file in rank zero mode (only on one process in multi-GPU setup)."""
     with open(path, "w+") as file:
         file.write(content)
@@ -129,7 +129,7 @@ def instantiate_loggers(logger_cfg: DictConfig) -> List[LightningLoggerBase]:
 
 
 @rank_zero_only
-def log_hyperparameters(object_dict: Dict[str, Any]) -> None:
+def log_hyperparameters(object_dict: dict) -> None:
     """Controls which config parts are saved by lightning loggers.
 
     Additionally saves: