提交 2215f46a 编写于 作者: W wuzewu

Fix the problem that hub.Module load fails in the multi-card environment

上级 a55bd6dd
...@@ -202,6 +202,12 @@ class Module(object): ...@@ -202,6 +202,12 @@ class Module(object):
else: else:
user_module_cls.source = '' user_module_cls.source = ''
user_module_cls.branch = '' user_module_cls.branch = ''
# In the case of multiple cards, the following code can set each process to use the correct place.
if issubclass(user_module_cls, paddle.nn.Layer):
place = paddle.get_device().split(':')[0]
paddle.set_device(place)
return user_module_cls return user_module_cls
@classmethod @classmethod
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册