Fix flaky test (#40175)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40175
Check that there is an increasing memory usage in the test
Test Plan: CI
Differential Revision: D22098192
Pulled By: ilia-cher
fbshipit-source-id: bbdbc71f66baf18514332a98d8927441c61ebc16
