I think your input should require gradients as it’s most likely a result of some other model operation.
If so, I assume you don’t need to set required_grad=True for layer.
Could you post some random inputs for all tensors, so that I could debug it?
Also, I’m not completely sure how your code works, but maybe the nested loop might be avoided using matrix operations.