Was able to stabilize the leak by setting OMP_NUM_THREADS=4. More info is here: https://github.com/pytorch/pytorch/issues/32008
A fix has been merged and looks to be available in future version 1.6. The nightly build works for this similar case: https://github.com/pytorch/pytorch/issues/40973`