fixed types

model/src/utils/load_model.py

@@ -4,10 +4,9 @@ import logging
 from pathlib import Path
 
 import torch
-from minio import Minio
 
 from model.src.models import VisualCommunicationModel
-from shared.datastore import get_model
+from shared.datastore import Datastore
 
 from .get_model_name import get_model_name
 
@@ -15,11 +14,11 @@ DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 def load_model(
-    client: Minio,
+    datastore: Datastore,
 ) -> VisualCommunicationModel:
     """Instantiate model with weights loaded from latest model saved in
     MinIO."""
-    assert isinstance(client, Minio)
+    assert isinstance(datastore, Datastore)
     # instantiate model
     model = VisualCommunicationModel()
     # get model object name
@@ -27,8 +26,7 @@ def load_model(
     model_object_name = get_model_name(path=model_name_path)
     logging.info('using model: %s', model_object_name)
     # load model from minio
-    model_checkpoint = get_model(
-        client=client,
+    model_checkpoint = datastore.get_model(
         object_name=model_object_name,
     )
     model.load_state_dict(model_checkpoint)