Fix memory leak in TraceEnumELBO

[fehiepsi]

Fixes #3068 and fixes #3014. I'm not sure what changes in PyTorch caused the issue but replacing x[...] with x.clone() fixes the memory leak.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[fritzo]

Great sleuthing! How did you find the offending line? Did you log tensor count or use any detective tooling code?

[fehiepsi]

How did you find the offending line?

Hi Fritz, I used your trick here to check for memory leaks in this example. After removing some gpu codes, I found that the leak also happens in CPU. Then I first copied the TraceEnum_ELBO implementation and removed line by line to see where is the offending code. That leads to Dice.compute_expectations and then this line with root._pyro_backward() call. When I print out the root, and root._pyro_backward I found that they are pretty simple expressions that only invoke _EinsumBackward and _LeafBackward, which are independent of the rest of pyro codes. I tried to play a little bit with _LeafBackward but couldn't replicate the leak so I guessed that the leak happened in some tensor operator of torch_marginal._EinsumBackward. It is lucky that only pyro.ops.einsum.torch_log.einsum is called and inside it, only operands[0][...] is called, so I found the leak. This process took me several hours though. (-:

[fritzo]

Thanks for your debugging story!

[qinqian]

Thank you @fehiepsi for the detailed fixing history!