Add output_mask (#66068)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/66068
Test Plan: Imported from OSS
Reviewed By: dagitses
Differential Revision: D31431802
Pulled By: albanD
fbshipit-source-id: 322aae5614dacb06fd45e513465b7a5cc11f4dbb
