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) 

残念ながら、これは動作していないようで、fooVにはuse_gpuがTrueのときに累積テンソルが含まれていません。

Answer 1

代替はまた

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

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

、GPUテンソルに変換するためにCPUテンソルにtype方法を使用することで、あなたは議論hereを見つけることができます。