Update CUSTOMRESNET.py

CUSTOMRESNET.py

@@ -2,11 +2,9 @@ import torch.nn.functional as F
 import torch
 import torch.nn as nn
 
-class Net(nn.Module):
+class MyModel(nn.Module):
     def __init__(self):
-        super(Net, self).__init__()
-        #drop=0.01
-        # Preparation Layer
+        super().__init__()
         self.conv1 = nn.Sequential (
             nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1,bias=False),
             nn.BatchNorm2d(64),
@@ -57,8 +55,8 @@ class Net(nn.Module):
         # Fully connected
         self.fc = nn.Linear(512, 10, bias=True)
 
-
     def forward(self, x):
+        #x = x.unsqueeze(0)
         x = self.conv1(x)
 
         x = self.conv11(x)
@@ -77,18 +75,16 @@ class Net(nn.Module):
         #x = x.randn(512, 1)
 
 # squeeze the tensor to size 512x
-        x = x.squeeze(dim=[2, 3])
+        #x = x.squeeze(dim=[2, 3])
+        x = x.squeeze(dim=2)
+        x = x.squeeze(dim=2)
+        #x = x.squeeze(dim=2).squeeze(dim=3)
+        #x = x.squeeze(dim=2)
 
         #x = x.view(512, 10)
 
         x = self.fc(x)
 
         x = x.view(-1, 10)
+        x = F.log_softmax(x, dim=-1)
         return x
-        #y = F.log_softmax(x, dim=-1)
-        #return  y
-
-def model_summary(model,input_size):
-    model = Net().to(device)
-    summary(model, input_size=(3, 32, 32))
-    return model,input_size