RuntimeError: CUDA error (10): invalid device ordinal

ゝ一世哀愁。 2022-03-09 14:17 428阅读 0赞

两张卡跑的train,在本机一张卡上test不了
结果很简单,直接将 model.module 拿出来就行了
因为之前双卡用的parallel

  1. def __init__(self, module, device_ids=None, output_device=None, dim=0):
  2. super(DataParallel, self).__init__()
  3. if not torch.cuda.is_available():
  4. self.module = module
  5. self.device_ids = []
  6. return
  7. if device_ids is None:
  8. device_ids = list(range(torch.cuda.device_count()))
  9. if output_device is None:
  10. output_device = device_ids[0]

详见:parallel源码
https://pytorch.org/docs/stable/_modules/torch/nn/parallel/data_parallel.html#DataParallel

发表评论

表情:
评论列表 (有 0 条评论,428人围观)

还没有评论,来说两句吧...

相关阅读