美文网首页
[PyTorch] Move to cuda移动到GPU上执行

[PyTorch] Move to cuda移动到GPU上执行

作者: VanJordan | 来源:发表于2019-04-28 17:27 被阅读0次
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)
  • 对于每一个样例,都移动到cuda
 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