bug fix

Berrlinn · Dec 24, 2020 · fb992ac · fb992ac
1 parent 8eeffa0
commit fb992ac
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/vedacore/misc/checkpoint.py b/vedacore/misc/checkpoint.py
@@ -166,7 +166,9 @@ def optimizer_to_cpu(state_dict):
     for key, val in state_dict.items():
         tmp = dict()
         for k, v in val.items():
-            tmp[k] = v.cpu()
+            if isinstance(v, torch.Tensor):
+                v = v.cpu()
+            tmp[k] = v
         state_dict_cpu[key] = tmp
     return state_dict_cpu