updated to use new io functions

tests/model_io_test.py

@@ -1,37 +1,34 @@
 """Test that a model can be saved and loaded again."""
 
-from io import BytesIO
+from pathlib import Path
 
 import torch
+from dotenv import load_dotenv
 from torchinfo import summary
 
 from model.src.models import VisualCommunicationModel
-from shared.data_store import connect, put
-from shared.utils import check_env, setup_logging
+from shared.data_store import connect, put_model
+from shared.utils import setup_logging
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 if __name__ == '__main__':
-    # check that environment variables are set
-    check_env()
+    # load in env file
+    env_path = Path(__file__).parent.parent / 'server.env'
+    assert env_path.exists()
+    load_dotenv(env_path)
     # setup logging
     setup_logging()
+    # connect to minio
+    client = connect()
     # instantiate model
     model = VisualCommunicationModel().to(DEVICE)
     # show model weights
     summary(model)
-    # for param in model.parameters():
-    #     logging.info(param.data)
-    # save model to buffer
-    buffer = BytesIO()
-    torch.save(model.state_dict(), buffer)
-    print(f"buffer size: {len(buffer.getvalue())}")
-    # connect to minio
-    client = connect(bucket='models')
     # put buffer in minio bucket
-    hash_str = put(
+    hash_str = put_model(
         client=client,
-        buffer=buffer,
+        model=model,
     )
     print(f"hash string: {hash_str}")
     print('saved to Minio')