Minor refactor to normalize assignments (#45671)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/45671
This is a follow up on D23977080 (https://github.com/pytorch/pytorch/commit/2596113a79add9ccc6921829e2ba674de0c32c85) and https://github.com/pytorch/pytorch/pull/45474.
Test Plan: See D23977080 (https://github.com/pytorch/pytorch/commit/2596113a79add9ccc6921829e2ba674de0c32c85).
Reviewed By: z-a-f
Differential Revision: D24043125
fbshipit-source-id: 0c05930668533bfd7145fa605f3785484391130b
