Question

如果激活了标记use_gpu，我尝试自动设置变量以使用cuda。

通常我会这样做：

import torch
from torch import autograd
use_gpu = torch.cuda.is_available()

foo = torch.randn(5)

if use_gpu:
    fooV = autograd.Variable(foo.cuda())
else:
    fooV = autograd.Variable(foo)

但是，当我必须在代码的不同部分初始化变量时，为了更快地做事，我想做这样的事情：

import torch
from torch import autograd
use_gpu = torch.cuda.is_available()

class Variable(autograd.Variable):
    def __init__(self, data, *args, **kwargs):
        if use_gpu:
            data = data.cuda()
        super(Variable, self).__init__(data, *args, **kwargs)

foo = torch.randn(5)
fooV = Variable(foo)

不幸的是，当fooV为use_gpu时，这似乎无法正常工作且True不包含cuda Tensor。

Answer 1

另一种方法是在CPU Tensor上使用type方法将其转换为GPU Tensor，

if use_cuda:
    dtype = torch.cuda.FloatTensor
else:
    dtype = torch.FloatTensor

x = torch.rand(2,2).type(dtype)

另外，您可以找到讨论here。

在pytorch中自动设置cuda变量

1 个答案: