Could you try to call .cuda() on the tensor instead of the nn.Parameter?
Most likely custom_param isn’t a leaf variable anymore as it was created by .cuda().
Have a look at this post for more information.
1 Like