diff --git a/paddlehub/module/module.py b/paddlehub/module/module.py index d2729530abe40349ce617ea7297b4368bcdf22f0..e0a17046e0958f6778f11ac1f0b605df7bd49ab0 100644 --- a/paddlehub/module/module.py +++ b/paddlehub/module/module.py @@ -202,6 +202,12 @@ class Module(object): else: user_module_cls.source = '' 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 @classmethod