[PyTorch] Move to cuda移动到GPU上执行
def move_to_cuda(sample):
if len(sample) == 0:
return {}
def _move_to_cuda(maybe_tensor):
if torch.is_tensor(maybe_tensor):
return maybe_tensor.cuda()
elif isinstance(maybe_tensor, dict):
return {
key: _move_to_cuda(value)
for key, value in maybe_tensor.items()
}
elif isinstance(maybe_tensor, list):
return [_move_to_cuda(x) for x in maybe_tensor]
else:
return maybe_tensor
return _move_to_cuda(sample)
for sample in t:
sample = utils.move_to_cuda(sample) if use_cuda else sample
本文标题:[PyTorch] Move to cuda移动到GPU上执行
本文链接:https://www.haomeiwen.com/subject/wqejnqtx.html
网友评论