Hi,
I had a similar problem, but for me it was a customized backward step, when calling the standard convolution in the forward pass. I built a workaround by calling the convolution from torch.nn._functions, as you can see here: Call backward on function inside a backpropagation step
Maybe this could help you? Maybe anyone in here has some ideas regarding my problem?