From 9a107c7bb700c83cd96034ab1cfec7ac1008286d Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Sat, 13 Jul 2024 22:38:23 +0200
Subject: [PATCH] added model instantiation test

---
 tests/model_io_test.py | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)
 create mode 100644 tests/model_io_test.py

diff --git a/tests/model_io_test.py b/tests/model_io_test.py
new file mode 100644
index 0000000..9eccc96
--- /dev/null
+++ b/tests/model_io_test.py
@@ -0,0 +1,37 @@
+"""Test that a model can be saved and loaded again."""
+
+from io import BytesIO
+
+import torch
+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
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if __name__ == '__main__':
+    # check that environment variables are set
+    check_env()
+    # setup logging
+    setup_logging()
+    # 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(
+        client=client,
+        buffer=buffer,
+    )
+    print(f"hash string: {hash_str}")
+    print('saved to Minio')