I can’t find any obvious issue in your code. Could you try to use some random tensors and post an executable code snippet, which would reproduce this issue, please?
Also, which PyTorch version are you using?