Fix data_prefetcher test error (#774)

Summary: Pull Request resolved: #774

Reviewed By: galrotem

Differential Revision: D55514862

fbshipit-source-id: 43c0bb7d549dfeba4e57efb3c9b29dd025efe5cb

Author: diego-urgell

tests/utils/data/test_data_prefetcher.py

@@ -9,6 +9,7 @@
 
 import unittest
 from typing import Tuple
+from unittest.mock import MagicMock, patch
 
 import torch
 from torch.utils.data.dataset import Dataset, TensorDataset
@@ -36,3 +37,30 @@ def test_device_data_prefetcher(self) -> None:
         num_prefetch_batches = 2
         with self.assertRaisesRegex(ValueError, "expects a CUDA device"):
             _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches)
+
+    @patch("torch.cuda.Stream")
+    def test_num_prefetch_batches_data_prefetcher(self, mock_stream: MagicMock) -> None:
+        device = torch.device("cuda:0")
+
+        num_samples = 12
+        batch_size = 4
+        dataloader = torch.utils.data.DataLoader(
+            self._generate_dataset(num_samples, 2), batch_size=batch_size
+        )
+
+        with self.assertRaisesRegex(
+            ValueError, "`num_prefetch_batches` must be greater than 0"
+        ):
+            _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=-1)
+
+        with self.assertRaisesRegex(
+            ValueError, "`num_prefetch_batches` must be greater than 0"
+        ):
+            _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=0)
+
+        # no exceptions raised
+        _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=1)
+        _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=2)
+
+        # Check that CUDA streams were created
+        self.assertEqual(mock_stream.call_count, 2)

tests/utils/data/test_data_prefetcher_gpu.py

@@ -25,30 +25,6 @@ def _generate_dataset(self, num_samples: int, input_dim: int) -> Dataset[Batch]:
         labels = torch.randint(low=0, high=2, size=(num_samples,))
         return TensorDataset(data, labels)
 
-    @skip_if_not_gpu
-    def test_num_prefetch_batches_data_prefetcher(self) -> None:
-        device = torch.device("cuda:0")
-
-        num_samples = 12
-        batch_size = 4
-        dataloader = torch.utils.data.DataLoader(
-            self._generate_dataset(num_samples, 2), batch_size=batch_size
-        )
-
-        with self.assertRaisesRegex(
-            ValueError, "`num_prefetch_batches` must be greater than 0"
-        ):
-            _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=-1)
-
-        with self.assertRaisesRegex(
-            ValueError, "`num_prefetch_batches` must be greater than 0"
-        ):
-            _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=0)
-
-        # no exceptions raised
-        _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=1)
-        _ = CudaDataPrefetcher(dataloader, device, num_prefetch_batches=2)
-
     @skip_if_not_gpu
     def test_cuda_data_prefetcher(self) -> None:
         device = torch.device("cuda:0")