Fixing and adding docstrings

main.py

@@ -80,11 +80,23 @@ def main():
     parser.add_argument(
         "--metric",
         type=str,
-        default=["entropy"],
+        default="entropy",
         choices=["entropy", "f1", "recall", "precision", "accuracy"],
         nargs="+",
         help="Which metric to use for evaluation",
     )
+    parser.add_argument('--imagesize',
+                        type=int,
+                        default=28,
+                        help='Size of images')
+    parser.add_argument('--imagechannels',
+                        type=int,
+                        default=1,
+                        choices=[1,3],
+                        help='Number of color channels in the image.')
+
+
+
 
     # Training specific values
     parser.add_argument(

utils/metrics/EntropyPred.py

@@ -5,5 +5,8 @@ class EntropyPrediction(nn.Module):
     def __init__(self):
         super().__init__()
 
-    def __call__(self, y_true, y_false):
+    def __call__(self, y_true, y_false_logits):
         return
+
+    def __reset__(self):
+        pass

utils/models/magnus_model.py

@@ -3,15 +3,33 @@
 
 class MagnusModel(nn.Module):
     def __init__(self,
-                 imgsize: int,
-                 channels: int,
+                 imagesize: int,
+                 imagechannels: int,
                  n_classes:int=10):
+
+        """
+        Magnus model contains the model for Magnus' part of the homeexam. 
+        This class contains a neural network consisting of three linear layers of 133 neurons each,
+        with ReLU activation between each layer.
+
+        Args
+        ----
+            imagesize (int): Expected size of input image. This is needed to scale first layer input
+            imagechannels (int): Expected number of image channels. This is needed to scale first layer input
+            n_classes (int): Number of classes we are to provide. 
+
+        Returns
+        -------
+            MagnusModel (nn.Module): Neural network as described above in this docstring.
+        """
+
+
         super().__init__()
-        self.imgsize = imgsize 
-        self.channels = channels
+        self.imagesize = imagesize 
+        self.imagechannels = imagechannels
 
         self.layer1 = nn.Sequential(*([
-            nn.Linear(self.channels*self.imgsize*self.imgsize, 133),
+            nn.Linear(self.imagechannels*self.imagesize*self.imagesize, 133),
             nn.ReLU()
         ]))
         self.layer2 = nn.Sequential(*([
@@ -24,12 +42,23 @@ def __init__(self,
         ]))
 
     def forward(self, x):
+        """
+        Forward pass of MagnusModel
+
+        Args
+        ----
+            x (th.Tensor): Four-dimensional tensor in the form (Batch Size x Channels x Image Height x Image Width)
+                
+        Returns
+        -------
+            out (th.Tensor): Class-logits of network given input x
+        """
         assert len(x.size) == 4
 
         x = x.view(x.size(0), -1)
 
         x = self.layer1(x)
         x = self.layer2(x)
-        x = self.layer3(x)
+        out = self.layer3(x)
 
-        return x
+        return out