add unit test

mosaicml · Jun 13, 2024 · 2098118 · 2098118
1 parent f46f6ff
commit 2098118
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tests/data/test_packing.py b/tests/data/test_packing.py
@@ -206,6 +206,12 @@ def test_auto_packing_with_streaming_dataloader(tmp_path: Path):
         if batch_ix >= 3:
             break
 
+    assert loader.dataset.packing_ratio is not None
+    assert loader.dataset.packing_ratio == int(loader.batch_size / 6)
+
+    state_dict = loader.dataset.state_dict(num_samples=2, from_beginning=False)
+    assert state_dict['sample_in_epoch'] == 2 * loader.dataset.packing_ratio
+
 
 @pytest.mark.parametrize('packing_ratio', ['auto', 2.0])
 @patch(