You should register the trainable tensors as nn.Parameters, which will then make sure to automatically push them to the specified device.
Change:
self.train_w = torch.randn(420,requires_grad = True)

to
self.train_w = nn.Parameter(torch.randn(420))