I’m doing a similar approach in calculating this, however I’m getting OOM after a few batches. Suppose that you feed your data batch-by-batch to this code and get the HVP for each batch, separately. Any idea how to avoid the OOM?
I’m sure the reason for OOM is the “create_graph=True” option in the first grad calculation (which of course is necessary for the second grad calculation), thought I’d argue that during the second call for grad function you don’t need the create_graph=True option.
Basically, I’d like to free the computed graph after each HVP calculation.