I am trying to fix: “RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!” from a torchscript module.
Is there a convenient way to move a whole module onto a particular device? I’ve tried m.to(torch.device('cuda'))
and m.cuda()
Here is a minimal (not quite working) example:
class TestModule(nn.Module):
def __init__(self, net):
super(TestModule, self).__init__()
self.net = net
self.a = torch.rand(2)
def forward(self, x):
s = self.net(x)
return self.a * s
m = TestModule(net)
Initializing with a network on the cpu:
m.net[-1].weights
→ cpu
m.a.device
→ device(type='cpu')
Then:
m = m.to(torch.device('cuda'))
m.net[-1].weight
→ device='cuda:0'
m.a.device
→ device(type='cpu'))
Thanks!