suppress errors in torchsnapshot

Differential Revision: D56154094 fbshipit-source-id: 875fdff3a7d3f30493936c2a04fbed98700f93ac
pytorch · Apr 15, 2024 · 706ee13 · 706ee13
1 parent 8bbf219
commit 706ee13
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 0 deletions.
diff --git a/tests/gpu_tests/test_dtensor_io_preparer.py b/tests/gpu_tests/test_dtensor_io_preparer.py
@@ -72,6 +72,7 @@ async def test_dtensor_io_preparer(
         """
         Verify the basic behavior of DTensorIOPreparer prepare_write.
         """
+        # pyre-fixme[6]: For 2nd argument expected `Union[_SupportsArray[typing.Any],...
         device_mesh = DeviceMesh("cuda", mesh=mesh)
 
         if len(placements) > device_mesh.ndim:

diff --git a/tests/gpu_tests/test_dtensor_utils.py b/tests/gpu_tests/test_dtensor_utils.py
@@ -31,6 +31,7 @@ class TestDTensorUtils(DTensorTestBase):
     @skip_if_lt_x_gpu(WORLD_SIZE)
     # pyre-fixme[3]: Return type must be annotated.
     def test_is_sharded_is_replicated(self):
+        # pyre-fixme[6]: For 2nd argument expected `Union[_SupportsArray[typing.Any],...
         mesh = DeviceMesh("cuda", mesh=[[0, 1], [2, 3]])
         placements = [Replicate(), Shard(0)]
         local_tensor = torch.rand((16, 16))

diff --git a/tests/test_batcher.py b/tests/test_batcher.py
@@ -165,8 +165,10 @@ def dtensor_test_cases(
     dsts = []
     for idx in range(NUM_TENSORS):
         mesh = (
+            # pyre-fixme[6]: For 2nd argument expected `Union[_SupportsArray[typing.A...
             DeviceMesh("cuda", mesh=[0])
             if use_gpu and idx % 2 == 0
+            # pyre-fixme[6]: For 2nd argument expected `Union[_SupportsArray[typing.A...
             else DeviceMesh("cpu", mesh=[0])
         )
         srcs.append(

diff --git a/tests/test_test_utils.py b/tests/test_test_utils.py
@@ -123,6 +123,7 @@ def _create_dtensor() -> DTensor:
 
         local_tensor = torch.rand((dim_0, dim_1))
 
+        # pyre-fixme[6]: For 2nd argument expected `Union[_SupportsArray[typing.Any],...
         mesh = DeviceMesh("cpu", mesh=[[0, 1], [2, 3]])
         placements = [Replicate(), Shard(0)]
         dtensor = distribute_tensor(

diff --git a/torchsnapshot/test_utils.py b/torchsnapshot/test_utils.py
@@ -317,6 +317,10 @@ def _dtensor_test_case(
     replicated: bool,
 ) -> Tuple[DTensor, Entry, List[WriteReq]]:
     # WORLD_SIZE needs to be at least 4
+    # pyre-fixme[6]: For 2nd argument expected `Union[_SupportsArray[typing.Any],
+    #  _NestedSequence[Union[bool, bytes, complex, float, int, str]],
+    #  _NestedSequence[_SupportsArray[typing.Any]], bool, bytes, complex, float, int,
+    #  str, Tensor]` but got `List[List[int]]`.
     mesh = DeviceMesh("cuda", mesh=[[0, 1], [2, 3]])
     placements = [Replicate(), Shard(0)]
     local_tensor = rand_tensor(shape, dtype=dtype)