Maybe there’s another way of achieving my underlying goal, but in order to create tensors of the appropriate cudaness in my modules, I’ve been doing something like:
class Foo(nn.Module):
def __init__(self):
super().__init__()
self.torch_constr = torch
def cuda(self):
super().cuda()
self.torch_constr = torch.cuda
def forward(self, x):
state = self.torch_constr.FloatTensor(...).zero_()
...
The problem I get is that the cuda method seems to only get called on the top-level class? Not sure that it’s called on the children? What are standard way(s) to handle this?
(edit: I’ve gone back to doing torch_constr = torch.cuda if x.is_cuda else torch for now)
