From 5a7e1aab91bbc99a3063efea1f3cde1968e3026f Mon Sep 17 00:00:00 2001 From: Mihir Patel Date: Tue, 30 Apr 2024 17:08:05 -0400 Subject: [PATCH] fix (#3234) --- tests/test_notebooks.py | 2 ++ tests/utils/object_store/test_gs_object_store.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/tests/test_notebooks.py b/tests/test_notebooks.py index 3a8fe50bc3..e63fa192a4 100644 --- a/tests/test_notebooks.py +++ b/tests/test_notebooks.py @@ -133,6 +133,8 @@ def test_notebook(notebook: str, device: str, s3_bucket: str): ) if notebook_name == 'checkpoint_autoresume': pytest.skip('MNIST dataset download is flaky') + if notebook_name == 'exporting_for_inference': + pytest.skip('MNIST dataset download is flaky') try: import boto3 diff --git a/tests/utils/object_store/test_gs_object_store.py b/tests/utils/object_store/test_gs_object_store.py index cd4d8fa03e..73006580eb 100644 --- a/tests/utils/object_store/test_gs_object_store.py +++ b/tests/utils/object_store/test_gs_object_store.py @@ -12,6 +12,7 @@ from composer.loggers import RemoteUploaderDownloader from composer.trainer import Trainer from composer.utils import GCSObjectStore +from composer.optim import DecoupledSGDW from tests.common import RandomClassificationDataset, SimpleModel @@ -28,8 +29,10 @@ def test_gs_object_store_integration_json_auth(expected_use_gcs_sdk_val=True, cl model = SimpleModel() train_dataset = RandomClassificationDataset() train_dataloader = DataLoader(dataset=train_dataset) + optimizer = DecoupledSGDW(model.parameters(), lr=1e-4) trainer_save = Trainer( model=model, + optimizer=optimizer, train_dataloader=train_dataloader, save_folder='gs://mosaicml-internal-integration-testing/checkpoints/{run_name}', save_filename='test-model.pt', @@ -48,6 +51,7 @@ def test_gs_object_store_integration_json_auth(expected_use_gcs_sdk_val=True, cl trainer_load = Trainer( model=model, + optimizer=optimizer, train_dataloader=train_dataloader, load_path=f'gs://mosaicml-internal-integration-testing/checkpoints/{run_name}/test-model.pt', max_duration='2ba',