Certain members of a torch module aren't moved to GPU even if model.to(device) is called

user_11790637 0
Viewed 1447

An mwe is as follows:

import torch
import torch.nn as nn

class model(nn.Module):
    def __init__(self):
        super(model,self).__init__()
        self.mat = torch.randn(2,2)

    def forward(self,x):
        print('self.mat.device is',self.mat.device)
        x = torch.mv(self.mat,x)
        return x

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
m = model()
m.to(device)

x = torch.tensor([2.,1.])
x = x.to(device)

m(x)

The output is

self.mat.device is cpu

and right after that comes

Traceback (most recent call last):
  File "Z:\cudatest.py", line 21, in <module>
    print(m(x))
  File "E:\Python37\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "Z:\cudatest.py", line 11, in forward
    x = torch.mv(self.mat,x)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_mv

The code works fine if I set device = torch.device('cpu'). It seems that the problem is model.mat is not moved to GPU even after m.to(device) is called. Why doesn't whis work? How can I fix this? Please note the following:

  1. Even though this particular example can be fixed by using self.mat = nn.Linear(2,2) and x = self.mat(x) instead, in my original program, I need a temporary tensor to store some data in forward() that is also used in some arithmetics. How can I construct such a tensor and send it to GPU when calling m.to(device)

  2. It is not known in advance whether the computer has GPU or not. Therefore, writing self.mat = self.mat.cuda() is not a good solution for my case.

pytorch python
1 Answers

pytorch apply Module's methods such as .cpu(), .cuda() and .to() only to sub-modules, parameters and buffers, but NOT to regular class members. pytorch has no way of knowing that self.mat, in your case, is an actual tensor that should be moved around.

Once you decide if your mat should be a parameter or a buffer, simply register it accordingly, e.g.

class model(nn.Module):
    def __init__(self):
        super(model,self).__init__()
        self.register_buffer(name='mat', tensor=torch.randn(2,2))
user_17144100
Related