python - pytorch中自动设置cuda变量

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)

不幸的是，当 use_gpu 为 True 时，这似乎不起作用并且 fooV 不包含 cuda 张量。

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 .