fix(nyz): fix policy set device bug

opendilab · Jun 1, 2022 · 1e0c4a1 · 1e0c4a1
1 parent 43f6f01
commit 1e0c4a1
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/ding/policy/base_policy.py b/ding/policy/base_policy.py
@@ -75,7 +75,6 @@ def __init__(
             if len(set(self._enable_field).intersection(set(['learn']))) > 0:
                 self._rank = get_rank() if self._cfg.learn.multi_gpu else 0
                 if self._cuda:
-                    torch.cuda.set_device(self._rank % torch.cuda.device_count())
                     model.cuda()
                 if self._cfg.learn.multi_gpu:
                     bp_update_sync = self._cfg.learn.get('bp_update_sync', True)
@@ -84,7 +83,6 @@ def __init__(
             else:
                 self._rank = 0
                 if self._cuda:
-                    torch.cuda.set_device(self._rank % torch.cuda.device_count())
                     model.cuda()
             self._model = model
             self._device = 'cuda:{}'.format(self._rank % torch.cuda.device_count()) if self._cuda else 'cpu'