Fixed fc1 parameter

models/modelOne.py

@@ -13,7 +13,7 @@ class modelOne(nn.Module) :
         self.conv2 = nn.Conv2d(6, 16, 5, padding=2)
         self.batchNorm2 = nn.BatchNorm2d(16)
 
-        self.fc1 = nn.Linear(16*64*64, 512)
+        self.fc1 = nn.Linear(63504, 512)
         self.dropout = nn.Dropout(0.5)
 
         self.fc2 = nn.Linear(512, 84)
@@ -23,6 +23,7 @@ class modelOne(nn.Module) :
         x = self.pool(F.relu(self.batchNorm1(self.conv1(x))))
         x = self.pool(F.relu(self.batchNorm2(self.conv2(x))))
         x = torch.flatten(x, 1)
+        print("Flattened size:", x.shape[1]) 
         x = self.dropout(x)
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))