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!