search

bayesiains / nflows

Normalizing flows in PyTorch
MIT License
851 stars 119 forks source link

Fix memory leak in batch norm #11

Closed mj-will closed 4 years ago

mj-will commented 4 years ago

This pull request fixes the memory leak described in the linked issue: https://github.com/bayesiains/nflows/issues/10

The fix is simply to detach the graph prior to computing the running meaning and variance for batch norm.

arturbekasov commented 4 years ago

Hi Michael.

Yep, this is indeed a bug. We shouldn't be adding running mean/variance computation to the graph. Thanks for providing the testing script in the linked issue: I've used it to reproduce the problem.

A nice and simple fix: looks good to me. I've tested memory usage with the fix applied: works as expected.

Thanks a lot for taking the time to report and fix this.