From c903713adfda0397d8c3994d651f1a4003079247 Mon Sep 17 00:00:00 2001 From: GreenGarnets Date: Mon, 20 Jan 2020 16:25:21 +0900 Subject: [PATCH] bug fix --- net/model.py | 3 +-- train.py | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/net/model.py b/net/model.py index 20f7d64..34b19a8 100644 --- 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 index 7efe274..94d973f 100644 --- 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()