fix (#3234)

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

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',