pytorch

d2123205 - Address review comments in https://github.com/pytorch/pytorch/pull/26272 (#26587)

Address review comments in https://github.com/pytorch/pytorch/pull/26272 (#26587) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/26587 - ghstack-source-id: 90557226 Test Plan: unit tests Differential Revision: D17515048 fbshipit-source-id: 3459ee80efec29080060ec29d67642d789dd8749