Solved by tom in post #4 
              
              
                .to is computation if it changes the device, so your parameter tensors will not be leaves. Replacing .to(device) with a device=device keyword argument in the factory functions (torch.normal and torch.zeros) should solve make it work on GPU as well. 
Best regards 
Thomas