don't check memory format for empty tensors

[ghstack-poisoned]
pytorch · May 18, 2024 · e02f02a · e02f02a
1 parent ed27879
commit e02f02a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_modules.py b/test/test_modules.py
@@ -663,10 +663,10 @@ def inner_check_out_mem_format(output):
                 d = output.dim()
                 if (d == 4 and ((input_mem_format == torch.channels_last)
                                 or (module_mem_format == torch.channels_last and module_memformat_affects_out))):
-                    self.assertTrue(output.is_contiguous(memory_format=torch.channels_last))
+                    self.assertTrue(output.numel() == 0 or output.is_contiguous(memory_format=torch.channels_last))
                 elif (d == 5 and ((input_mem_format == torch.channels_last_3d)
                                   or (module_mem_format == torch.channels_last_3d and module_memformat_affects_out))):
-                    self.assertTrue(output.is_contiguous(memory_format=torch.channels_last_3d))
+                    self.assertTrue(output.numel() == 0 or output.is_contiguous(memory_format=torch.channels_last_3d))
                 else:
                     self.assertTrue(output.is_contiguous())
             return self._traverse_obj(output, inner_check_out_mem_format)