Hi,
The problem here is that your implementation of the backward is not correct.

You don’t need to use torch.no_grad() is a custom Function
Why not run the exact same code (without the conversion to numpy) in your criterion_mse() function so that it will be autodiffed and you get backward for free?
If you don’t do that, you need to implement the backward yourself. And in particular, the error you see here is because the gradient your return for input don’t have the same size as `input.

PS: I updated your post to make the code more readable using triple backticks.