Skip to content

Commit 2e45e81

Browse files
committed
fix array dl tests gpu vs cpu (#139)
1 parent 53e3533 commit 2e45e81

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

tests/unit/dataloader/test_array_dataloader.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,4 +51,7 @@ def test_dataloader(tmpdir, dataset, cpu, num_rows):
5151
for col in columns:
5252
start = idx * int(num_rows / 10)
5353
end = start + int(num_rows / 10)
54-
assert np.all(batch[0][col].ravel() == df[col].iloc[start:end])
54+
if cpu:
55+
assert np.all(batch[0][col].ravel() == df[col].iloc[start:end].to_numpy())
56+
else:
57+
assert np.all(batch[0][col].ravel() == df[col].iloc[start:end])

0 commit comments

Comments
 (0)