Enable torch.cuda.nvtx typechecking (pytorch#43443)

Summary:
Add pyi file covering torch._C.nvtx submodule

Fixes pytorch#43436

Pull Request resolved: pytorch#43443

Reviewed By: ezyang

Differential Revision: D23280188

Pulled By: malfet

fbshipit-source-id: 882860cce9feb0b5307c8b7c887f4a2f2c1548a2

Author: malfet

mypy.ini

@@ -204,9 +204,6 @@ ignore_errors = True
 [mypy-torch.cuda.nccl]
 ignore_errors = True
 
-[mypy-torch.cuda.nvtx]
-ignore_errors = True
-
 [mypy-torch._lobpcg]
 ignore_errors = True
 

torch/_C/_nvtx.pyi

@@ -0,0 +1,4 @@
+# Defined in torch/csrc/cuda/shared/nvtx.cpp
+def rangePushA(message: str) -> int: ...
+def rangePop() -> int: ...
+def markA(message: str) -> None: ...

torch/cuda/nvtx.py

@@ -10,7 +10,7 @@ def _fail(*args, **kwargs):
         rangePop = _fail
         markA = _fail
 
-    _nvtx = _NVTXStub()
+    _nvtx = _NVTXStub()  # type: ignore[assignment]
 
 __all__ = ['range_push', 'range_pop', 'mark']