fixed mypy types

model/src/utils/get_class.py

@@ -1,8 +1,9 @@
 from importlib import import_module
-from types import ModuleType
+
+from torch.nn import Module
 
 
-def get_class(path: str) -> ModuleType:
+def get_class(path: str) -> Module:
     parts = path.split('.')
     module_path = '.'.join(parts[:-1])
     class_name = parts[-1]

model/src/utils/load_model.py

@@ -5,8 +5,8 @@ from pathlib import Path
 
 import torch
 from minio import Minio
-from src.models import VisualCommunicationModel
 
+from model.src.models import VisualCommunicationModel
 from shared.data_store import get_model
 
 from .get_model_name import get_model_name

model/train.py

@@ -13,11 +13,11 @@ from ignite.handlers.param_scheduler import create_lr_scheduler_with_warmup
 from ignite.handlers.tensorboard_logger import TensorboardLogger
 from ignite.handlers.tqdm_logger import ProgressBar
 from ignite.metrics import Average, Loss, RunningAverage
-from src.utils import VCDADataset, get_class
 from torch import nn
 from torch.optim.lr_scheduler import ExponentialLR
 from torch.utils.data import DataLoader
 
+from model.src.utils import VCDADataset, get_class
 from shared.data_store import connect_minio
 
 
@@ -90,7 +90,7 @@ val_metrics = {
 }
 evaluator = create_supervised_evaluator(
     model=model,
-    metrics=val_metrics,
+    metrics=val_metrics,  # type: ignore
     device=device,
 )
 ProgressBar(desc='Val', ncols=80).attach(evaluator)