Fix asserts in tests (#72864)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72864
Fixes #72860
Test Plan: Imported from OSS
Reviewed By: rohan-varma
Differential Revision: D34246987
Pulled By: H-Huang
fbshipit-source-id: 1ba47585533aff4cff9beec49bdc801f8320ffc8
(cherry picked from commit 03e45ceb890d72216950a9c3d5cd648b02e6a557)
