CUDA OOM Error about captum

[vivekmig]

Hi @zwacke , it depends on what you've set for n_steps. Integrated Gradients needs to compute gradients at each of the n_steps points, so one input image expands to a batch of n_step evaluations, which is likely causing the out of memory issue.

You can try either reducing n_steps to a smaller value or otherwise using the internal_batch_size argument in IG, which splits the evaluations into batches with at most internal_batch_size examples. Setting internal_batch_size to a smaller value (corresponding to a batch size that can fit in memory) while keeping n_steps the same should also work. Let us know if this resolves the issue!

[zwacke]

Yes, thank you it did. Defaulting to a working internal_batch_size seems to be the recommendable approach. If I get it right, this especially helps when wrapping an attribution method in NoiseTunnel, or any method that introduces further expansion of the overall input batch size.

[tranvnhan]

I encountered the same problem today. By reducing n_steps to a smaller value, it works now.

[NarineK]

Hi @zwacke, did Vivek's suggestion help ? If so can we close the issue ?

[NarineK]

Nice! Glad that it helped. Feel free to close the issue if it's fixed!