Add annotations to torchrec/tensor_types.py (#1648)

Summary: Pull Request resolved: #1648 Reviewed By: azad-meta Differential Revision: D52957237 fbshipit-source-id: 7515239cfc5a8a5a865d11cda6ae39768f0acfae
pytorch · Jan 22, 2024 · 239c033 · 239c033
1 parent ece441d
commit 239c033
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrec/tensor_types.py b/torchrec/tensor_types.py
@@ -72,7 +72,7 @@ class UIntXTensor(torch.Tensor):
     __torch_function__ = torch._C._disabled_torch_function_impl
 
     @staticmethod
-    def __new__(cls, N, elem):
+    def __new__(cls, N: int, elem):
         assert elem.dtype is torch.uint8
         # pyre-ignore
         return torch.Tensor._make_wrapper_subclass(