From fb992aca6ce1bb024d07aef669c349e3ecc6516c Mon Sep 17 00:00:00 2001
From: "hongxiang.cai" <hongxiang.cai@media-smart.cn>
Date: Thu, 24 Dec 2020 21:43:37 +0800
Subject: [PATCH] bug fix

---
 vedacore/misc/checkpoint.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/vedacore/misc/checkpoint.py b/vedacore/misc/checkpoint.py
index e189417..ce43496 100644
--- 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