Update train.py

解决模型评估过程出现显存爆炸
ice-tong · Mar 14, 2020 · 0ef099d · 0ef099d
1 parent 0df12e1
commit 0ef099d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/train.py b/train.py
@@ -80,8 +80,8 @@ def train():
             optimizer.step()
 
             acc = calculat_acc(output, target)
-            acc_history.append(acc)
-            loss_history.append(loss)
+            acc_history.append(float(acc))
+            loss_history.append(float(loss))
         print('train_loss: {:.4}|train_acc: {:.4}'.format(
                 torch.mean(torch.Tensor(loss_history)),
                 torch.mean(torch.Tensor(acc_history)),
@@ -99,7 +99,7 @@ def train():
             output = cnn(img)
 
             acc = calculat_acc(output, target)
-            acc_history.append(acc)
+            acc_history.append(float(acc))
             loss_history.append(float(loss))
         print('test_loss: {:.4}|test_acc: {:.4}'.format(
                 torch.mean(torch.Tensor(loss_history)),
@@ -110,4 +110,4 @@ def train():
 
 if __name__=="__main__":
     train()
-    pass
+    pass