Add more kernel launch checks (#53286)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/53286
Test Plan: Sandcastle
Reviewed By: ngimel
Differential Revision: D26818164
fbshipit-source-id: 01ba50dc7e4a863e26c289d746bc5b95aa76d3cc
