init commit

1 files changed, 11 insertions, 0 deletions

diff --git a/files/tests/test_data_io.py b/files/tests/test_data_io.py
new file mode 100644
index 0000000..1f2ccd8
--- /dev/null
+++ b/files/tests/test_data_io.py
@@ -0,0 +1,11 @@
+import torch
+from files.data_io.dataset_loader import get_dataloader
+
+
+def test_dataloader_shape():
+    """Smoke test: verify dataloader output shape."""
+    train_loader, _ = get_dataloader("data_io/configs/shd.yaml")
+    x, y = next(iter(train_loader))
+    assert isinstance(x, torch.Tensor)
+    assert x.ndim == 3
+    assert y.ndim == 1