Hi,
For the input, you can do the same thing: input = torch.randn(128, 20, device="cuda") But for the nn.Module you will have to first create it and then send it to the GPU:
m = nn.Linear(20, 30)
m.cuda()