fix forward test_memory_planning.py (#119109)
Summary: fixes a broken test, also makes it run in fbcode correctly
Test Plan: test
Reviewed By: angelayi
Differential Revision: D53373709
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119109
Approved by: https://github.com/angelayi
