Hi,
Since the error shows that the failure happens in your custom Function. I think the simplest is to print the size directly in your manual backward.
In general, for the backward, the grad you get as input will have the same size as the output of the forward. And you must return Tensors with the same size as the input of the forward.