Hi,
Thanks for the nice run-able code!
When i run it, it complains that the size of the gradients returned by your custom functions are wrong. Do you see the same issue?
I guess you need to keep these values in the einsum for the Hadamard function, but you can change the return statement to ..., grad_weight.sum(dim=(-1, -2)), grad_bias.sum(dim=(-1, -2)),... to fix that.