Almost, as I don’t think directly assigning tensors to ctx and unpacking them is a valid approach since it could cause memory leaks, if I’m not mistaken.
Use ctx.save_for_backward and ctx.saved_tensors instead as shown in the tutorial.
1 Like