Added test_load_model

Author: Seilmast

tests/test_metrics.py

@@ -3,7 +3,13 @@
 import pytest
 
 from CollaborativeCoding.load_metric import MetricWrapper
-from CollaborativeCoding.metrics import Accuracy, F1Score, Precision, Recall
+from CollaborativeCoding.metrics import (
+    Accuracy,
+    EntropyPrediction,
+    F1Score,
+    Precision,
+    Recall,
+)
 
 
 @pytest.mark.parametrize(
@@ -17,7 +23,7 @@
         ("accuracy", randint(2, 10), True),
         ("precision", randint(2, 10), False),
         ("precision", randint(2, 10), True),
-        # TODO: Add test for EntropyPrediction
+        ("EntropyPrediction", randint(2, 10), False),
     ],
 )
 def test_metric_wrapper(metric, num_classes, macro_averaging):

tests/test_wrappers.py

@@ -1,30 +1,30 @@
 from CollaborativeCoding import load_data, load_metric, load_model
 
-# def test_load_model():
-#     import torch as th
-
-#     image_shape = (1, 16, 16)
-#     num_classes = 4
-
-#     dummy_img = th.rand((1, *image_shape))
-
-#     modelnames = [
-#         "magnusmodel",
-#         "christianmodel",
-#         "janmodel",
-#         "solveigmodel",
-#         "johanmodel",
-#     ]
-
-#     for name in modelnames:
-#         print(name)
-#         model = load_model(name, image_shape=image_shape, num_classes=num_classes)
-
-#         with th.no_grad():
-#             output = model(dummy_img)
-#             assert output.size() == (1, 4), (
-#                 f"Model {name} returned image of size {output}. Expected (1,4)"
-#             )
+
+def test_load_model():
+    import torch as th
+
+    image_shape = (1, 16, 16)
+    num_classes = 4
+
+    dummy_img = th.rand((1, *image_shape))
+
+    modelnames = [
+        "magnusmodel",
+        "christianmodel",
+        "janmodel",
+        "solveigmodel",
+        "johanmodel",
+    ]
+
+    for name in modelnames:
+        model = load_model(name, image_shape=image_shape, num_classes=num_classes)
+
+        with th.no_grad():
+            output = model(dummy_img)
+            assert output.size() == (1, 4), (
+                f"Model {name} returned image of size {output}. Expected (1,4)"
+            )
 
 
 def test_load_data():