pytorch instance tensor not moved to gpu even with explicit cuda() call

user_6471140 0
Viewed 5268

I'm working on a project where the model requires access to a tensor that i declare in the constructor init of the class (im sub-classing torch.nn.Module class) and then i need to use this tensor in the forward() method via a simple matmul() , the model is sent to gpu via a cuda() call:

model = Model()
model.cuda()

However when i do forward-propagation of a simple input X through:

model(X) # or model.forward(X)

I get

RuntimeError: Expected object of type torch.cuda.FloatTensor but found type torch.FloatTensor for argument #2 'mat2'

Indicating that the second argument of matmul(the instance tensor i declared) is on CPU and it was expected on GPU (as the rest of the model and data).

In matmul, the tensor is transposed via matrix.t()

I even tried overriding the cuda() method thorugh:

def cuda(self):
    super().cuda()
    self.matrix.cuda()

The data is already in the GPU ,meaning the following line of code was already executed:

X = X.cuda()

Also the error explcitly says argument 2 of matmul which for this case is the tensor(called matrix) not X.

pytorch machine-learning artificial-intelligence python
2 Answers

Let's assume the following:

  1. X is moved correctly to the GPU

  2. The tensor declared in the Model class is a simple attribute.

    i.e. Something like the following:

class Model(nn.Module):
   def __init__(self):
       super().__init__()
       self.matrix = torch.randn(784, 10)
       
   def forward(self, x):
       return torch.matmul(x, self.matrix)

If so, your first attempt wouldn't work because the nn.Module.cuda() method only moves all of the Parameters and Buffers to the GPU.

You would need to make Model.matrix a Parameter instead of regular attribute. Wrap it in the parameter class. Something like:

self.matrix = nn.Parameter(torch.randn(784, 10))

Now, instead of automatically casting to the GPU like above, you tried to manually call the .cuda() method on Model.matrix within the override.

This doesn't work either because of a subtle difference between the nn.Module.cuda() method and the torch.Tensor.cuda() method.

While nn.Module.cuda() moves all the Parameters and Buffers of the Module to GPU and returns itself, torch.Tensor.cuda() only returns a copy of the tensor on the GPU.

The original tensor is unaffected.

In summary, either:

  1. Wrap your matrix attribute as a Parameter or
  2. Assign the GPU copy back to matrix via:
self.matrix = self.matrix.cuda()

In your override.

I would suggest the first.

user_69085140

I would like to highlight this from @Vaisakh's answer:

While nn.Module.cuda() moves all the Parameters and Buffers of the Module to GPU and returns itself, torch.Tensor.cuda() only returns a copy of the tensor on the GPU.

In other words, as @Umang_Gupta says in his comment:

# if m is a Module, you do:
m.cuda()

# if t is a Tensor, you do:
t = t.cuda()
user_4208670
Related