Unction CudnnConvolutionBackward returned an invalid gradient at index 1 - expected type torch.FloatTensor but got torch.cuda.FloatTensor

Could you print the device attribute of all parameters before calling train?

for name, param in model.named_parameters():
    device = param.device
    if device == torch.device('cpu'):
        print('ERROR in', name, device)