Hi,
You are using “old style” autograd.Function and this can happen.
Can you change your code to use the new style as described in this doc 22 (basically remove the init and pass arguments directly to the forward, make forward and backward static method, save stuff in ctx. during the forward if you need them in the backward) and see if it solves the issue.