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 张量上使用 type 方法将其转换为 GPU 张量，

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

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

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

pytorch中自动设置cuda变量

Automatically set cuda variables in pytorch

python

pytorch