fix save and load from non 0 device

Oneflow-Inc · Sep 22, 2023 · 905ccaf · 905ccaf
1 parent 4e60628
commit 905ccaf
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/oneflow/nn/graph/graph.py b/python/oneflow/nn/graph/graph.py
@@ -1069,7 +1069,8 @@ def _fill_sub_destination(dest_dict, name_list, tensor_tuple):
             assert len(tensor_tuple) == len(name_list)
             for name_idx in range(len(name_list)):
                 tensor_item = tensor_tuple[name_idx]
-                dest_dict[name_list[name_idx]] = (tensor_item, tensor_item.device.type)
+                device_str = ':'.join((tensor_item.device.type, str(tensor_item.device.index)))
+                dest_dict[name_list[name_idx]] = (tensor_item, device_str)
 
         # This is original outputs is needed to build output buffer.
         tuple_idx = -1