美文网首页
[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