Browse Source

修复一个data_parallel情况下的gpu不从0开始导致的bug

tags/v0.6.0
yh_cc 5 years ago
parent
commit
ed2309719c
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      fastNLP/core/utils.py

+ 1
- 1
fastNLP/core/utils.py View File

@@ -409,7 +409,7 @@ def _move_model_to_device(model, device):

if device is None:
if isinstance(model, torch.nn.DataParallel):
model.cuda()
model.cuda(model.device_ids[0])
return model
else:
if not torch.cuda.is_available() and ((isinstance(device, str) and device!='cpu') or


Loading…
Cancel
Save