From 022facb94fb607e3956f486f797672d56ec576d6 Mon Sep 17 00:00:00 2001 From: edknv Date: Thu, 23 Mar 2023 11:01:26 -0700 Subject: [PATCH] Change row lengths unit tests to use cpu and single gpu --- tests/unit/dataloader/test_torch_dataloader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/dataloader/test_torch_dataloader.py b/tests/unit/dataloader/test_torch_dataloader.py index 61c8960b..2b771da8 100644 --- a/tests/unit/dataloader/test_torch_dataloader.py +++ b/tests/unit/dataloader/test_torch_dataloader.py @@ -420,7 +420,7 @@ def test_offsets(): assert feature_tensor == feature_result -@pytest.mark.parametrize("device", ["cpu", 0, 1] if HAS_GPU else ["cpu"]) +@pytest.mark.parametrize("device", ["cpu", 0] if HAS_GPU else ["cpu"]) def test_row_lengths_to_offsets_device(device): dataset = Dataset(make_df({"a": [1]})) loader = torch_dataloader.Loader(dataset, batch_size=1)