Fix daily tests (mosaicml#2891)

* fix daily * revert --------- Co-authored-by: Daniel King <[email protected]>
mvpatel2000 · Jan 22, 2024 · c247bc3 · c247bc3
1 parent fbc4d85
commit c247bc3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py
@@ -80,6 +80,7 @@ def modify_cell_source(tb: TestbookNotebookClient, notebook_name: str, cell_sour
         cell_source = cell_source.replace('batch_size = 1024', 'batch_size = 64')
         cell_source = cell_source.replace('download=True', 'download=False')
     if notebook_name == 'auto_microbatching':
+        cell_source = cell_source.replace('batch_size = 2048', 'batch_size = 1024')
         cell_source = cell_source.replace('download=True', 'download=False')
     if notebook_name == 'migrate_from_ptl':
         cell_source = cell_source.replace('batch_size=256', 'batch_size=64')