Call .stop() on dataloader in test_dataloader_schema (#1848)

NVIDIA-Merlin · Jul 3, 2023 · f9a2891 · f9a2891
1 parent a775cff
commit f9a2891
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/tests/unit/loader/test_tf_dataloader.py b/tests/unit/loader/test_tf_dataloader.py
@@ -509,3 +509,5 @@ def test_dataloader_schema(tmpdir, df, dataset, batch_size, engine, device):
 
     num_label_cols = batch[1].shape[1] if len(batch[1].shape) > 1 else 1
     assert num_label_cols == len(label_name)
+
+    data_loader.stop()
diff --git a/tests/unit/loader/test_torch_dataloader.py b/tests/unit/loader/test_torch_dataloader.py
@@ -526,3 +526,5 @@ def test_dataloader_schema(tmpdir, df, dataset, batch_size, engine, device):
 
     num_label_cols = batch[1].shape[1] if len(batch[1].shape) > 1 else 1
     assert num_label_cols == len(label_name)
+
+    data_loader.stop()