The test starts failing flakily possibly after https://github.com/pytorch/pytorch/pull/143787 lands and needs to be updated cc @seemethere @malfet @pytorch/pytorch-dev-infra @mergennachin