Hi,
All the code outside the custom Function is autodiffed. And so gradient will be computed for each op you do. So all will work fine here!
Note that in the custom function, you should not save inputs in the context but use ctx.save_for_backward(y) and y, = ctx.saved_tensors to avoid memory leak