improve memory usage

yinboc · Nov 25, 2018 · 1ad07d8 · 1ad07d8
1 parent a3f8f1e
commit 1ad07d8
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/train.py b/train.py
@@ -88,6 +88,8 @@ def save_model(name):
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
+
+            proto = None; logits = None; loss = None
 
         tl = tl.item()
         ta = ta.item()
@@ -114,6 +116,8 @@ def save_model(name):
 
             vl.add(loss.item())
             va.add(acc)
+
+            proto = None; logits = None; loss = None
 
         vl = vl.item()
         va = va.item()