Fixes failure in test_dataloader.py that occurs on jetson boards (#64757)

puririshi98 · facebook-github-bot · commit 2ae938e15e69 · 2021-09-13T10:11:04.000-07:00
Summary: CUDA IPC is not supported for jetsons Pull Request resolved: #64757 Reviewed By: jbschlosser Differential Revision: D30900593 Pulled By: ejguan fbshipit-source-id: c6b2e8a9746276fdb4a009b6412e47cc8aac69f2
diff --git a/test/test_dataloader.py b/test/test_dataloader.py
@@ -67,7 +67,11 @@
 # as well during the execution of this test suite, and it will cause
 # CUDA OOM error on Windows.
 TEST_CUDA = torch.cuda.is_available()
-
+if TEST_CUDA:
+    dev_name = torch.cuda.get_device_name(torch.cuda.current_device()).lower()
+    IS_JETSON = 'xavier' in dev_name or 'nano' in dev_name or 'jetson' in dev_name or 'tegra' in dev_name
+else:
+    IS_JETSON = False
 
 if not NO_MULTIPROCESSING_SPAWN:
     # We want to use `spawn` if able because some of our tests check that the
@@ -1334,8 +1338,8 @@ def test_multiprocessing_contexts(self):
         counting_ds_n = 11
         dl_common_args = dict(num_workers=3, batch_size=3, pin_memory=(not TEST_CUDA))
         for ctx in supported_multiprocessing_contexts:
-            # windows doesn't support sharing cuda tensor; ROCm does not yet fully support IPC
-            if ctx in ['spawn', 'forkserver'] and TEST_CUDA and not IS_WINDOWS:
+            # windows and jetson devices don't support sharing cuda tensor; ROCm does not yet fully support IPC
+            if ctx in ['spawn', 'forkserver'] and TEST_CUDA and not IS_WINDOWS and not IS_JETSON:
                 ds_cls = CUDACountingDataset
             else:
                 ds_cls = CountingDataset
