updated johan_model to flatten the input in the forward loop

Author: Johanmkr

utils/models/johan_model.py

@@ -46,6 +46,7 @@ def __init__(self, image_shape, num_classes):
         self.relu = nn.ReLU()
 
     def forward(self, x):
+        x = x.flatten()
         for layer in [self.fc1, self.fc2, self.fc3, self.fc4]:
             x = layer(x)
             x = self.relu(x)