v2

2020-10-22 22:42:01 +08:00
parent 0019d4034c
commit 376f5caeb7
11 changed files with 140 additions and 29 deletions
--- a/engine/base/i2i.py
+++ b/engine/base/i2i.py
@@ -101,9 +101,12 @@ class EngineKernel(object):


 def _remove_no_grad_loss(loss_dict):
+    need_to_pop = []
    for k in loss_dict:
        if not isinstance(loss_dict[k], torch.Tensor):
-            loss_dict.pop(k)
+            need_to_pop.append(k)
+    for k in need_to_pop:
+        loss_dict.pop(k)
    return loss_dict