Followup of kron PR (#51045)
Summary:
Followup of https://github.com/pytorch/pytorch/pull/50927
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51045
Reviewed By: mruberry
Differential Revision: D26089204
Pulled By: ngimel
fbshipit-source-id: 77291dd83fba32d6f80a8540910b112a1d85a892
