fix finalize features floods CPU memory: remove log and load mapped gradients to device

MadryLab / trak

A fast, effective data attribution method for neural networks in PyTorch

https://trak.csail.mit.edu/

MIT License

175 stars 24 forks source link

Closed JD-ETH closed 9 months ago

JD-ETH commented 10 months ago

loading data directly on device instead of moving from CPU to GPU in score computation steps

kristian-georgiev commented 9 months ago

Great catch @JD-ETH, thanks!