relax size check in flatten_for_scatter_gather (#40573)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40573
Per title, to workaround apex sbn bug.
Test Plan: Covered by existing tests
Reviewed By: blefaudeux
Differential Revision: D22236942
fbshipit-source-id: ddb164ee347a7d472a206087e4dbd16aa9d72387
