Add batch_size to 1 if not provided for testing (#635)

mosaicml · Mar 21, 2024 · 6c6e143 · 6c6e143
1 parent d8bf491
commit 6c6e143
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/regression/utils.py b/regression/utils.py
@@ -90,6 +90,8 @@ def get_streaming_dataset_params(kwargs: dict[str, str]) -> dict[str, Any]:
         dataset_params['num_canonical_nodes'] = int(kwargs['num_canonical_nodes'])
     if 'batch_size' in kwargs:
         dataset_params['batch_size'] = int(kwargs['batch_size'])
+    else:
+        dataset_params['batch_size'] = 1
     if 'shuffle' in kwargs:
         dataset_params['shuffle'] = kwargs['shuffle'].lower().capitalize() == 'True'
     if 'shuffle_algo' in kwargs: