diff options
Diffstat (limited to 'files/tests/test_data_io.py')
| -rw-r--r-- | files/tests/test_data_io.py | 11 |
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 |
