bug fix

GreenGarnets · Jan 20, 2020 · c903713 · c903713
1 parent fd99fe5
commit c903713
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/net/model.py b/net/model.py
@@ -57,7 +57,6 @@ def __init__(self):
 
 
     def forward(self, x):
-        batch = x.size(0)
 
         out = self.conv1(x)
         out = self.conv2(out)
@@ -68,7 +67,7 @@ def forward(self, x):
         out = self.conv7(out)
         out = self.conv8(out)
 
-        out = out.reshape(batch,out.size(1)* out.size(2))
+        out = out.reshape(out.size(0), out.size(1)* out.size(2))
         #print(out.shape)
 
         out = self.fc(out)

diff --git a/train.py b/train.py
@@ -135,7 +135,7 @@ def main():
                 if i+batch < input.shape[0] :
                     optimizer.zero_grad()
 
-                    pred = model(input[i:i+batch],batch)
+                    pred = model(input[i:i+batch])
 
 
                     loss = criterion(pred.squeeze(), target[i:i+batch].squeeze())
@@ -173,7 +173,7 @@ def main():
             output = []
             for i in range(0,input.shape[0], batch):
                 if i+batch < input.shape[0] :
-                    pred = model(input[i:i+batch],batch)
+                    pred = model(input[i:i+batch])
                     pred = pred.to(torch.device("cpu"))
                     if i == 0 :                        
                         output = pred.tolist()