Thanks for the update. The output shows different shapes of the gradient and parameter so could you explain what the gradient would represent and how it should be applied on the parameter?
I don’t know why the GPU run isn’t failing, but assume it should and might be missing a check.
EDIT: I just checked this wrong assignment and it’s failing for CPU and CUDA tensors in my setup, so the check seems to be OK in a recent master build.