Move the check on orig_weight sizes. (#40200)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40200
Since original weights are removed by default in mobile build, the check must
be moved to a place where orig_weight is still valid.
Test Plan:
CI
Plus observed a model run crash which was resolved after this change.
Reviewed By: supriyar
Differential Revision: D22101562
fbshipit-source-id: 9543e69a415beaef2a9fb92dc9cd87d636174d51
