From 905ccaf23bc5bfd6d844ad290d188a942016fe12 Mon Sep 17 00:00:00 2001 From: strint Date: Fri, 22 Sep 2023 13:53:47 +0000 Subject: [PATCH] fix save and load from non 0 device --- python/oneflow/nn/graph/graph.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/oneflow/nn/graph/graph.py b/python/oneflow/nn/graph/graph.py index da5b8556c94..ab8b818d822 100644 --- 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