From b88e177305b5ea550ae2ce87c187a2aea07ae41b Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Date: Thu, 22 Jun 2023 16:58:15 +0100
Subject: [PATCH] Call `.stop()` on dataloader in `test_dataloader_schema`

---
 tests/unit/loader/test_tf_dataloader.py    | 2 ++
 tests/unit/loader/test_torch_dataloader.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/tests/unit/loader/test_tf_dataloader.py b/tests/unit/loader/test_tf_dataloader.py
index e1ba7bfe7c4..b51d0681903 100644
--- 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
index 2617fa0845d..99026e2ba7e 100644
--- 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()