Browse Source

修复_move_model_to_device中当device=torch.device('cpu')会触发的bug

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

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

@@ -223,8 +223,8 @@ def _move_model_to_device(model, device):
model.cuda()
return model
else:
if not torch.cuda.is_available() and (
device != 'cpu' or (isinstance(device, torch.device) and device.type != 'cpu')):
if not torch.cuda.is_available() and ((isinstance(device, str) and device!='cpu') or
(isinstance(device, torch.device) and device.type != 'cpu')):
raise ValueError("There is no usable gpu. set `device` as `cpu` or `None`.")

if isinstance(model, torch.nn.DataParallel):


Loading…
Cancel
Save