Hi,
Is there a reason why you’re not using the checkpoint 3 tool? It looks like exactly what you’re trying to reimplement.
For your code, the forward pass in a Function already runs in a torch.no_grad() block. So no need to add another one.
The fact that you return a tuple might be a problem. Can you just return your three tensors like return res_pred, add_out, factor ?