improve enforce in cross_entroy_op (#23062)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/23062
as title
Reviewed By: xianjiec, BIT-silence
Differential Revision: D16374601
fbshipit-source-id: 62219c6abde311ebc8a0e6a03cfb517d80bb52b5
