epochs to batches, also revert trying to make custom resnet

mosaicml · Nov 29, 2024 · dd39ec3 · dd39ec3
1 parent eafe347
commit dd39ec3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/tests/algorithms/algorithm_settings.py b/tests/algorithms/algorithm_settings.py
@@ -207,19 +207,19 @@
         'model': (
             composer_resnet,
             {
-                'model_name': 'resnet18',
+                'model_name': 'resnet50',
                 'num_classes': 2,
             },
         ),
         'dataset': (
             RandomImageDataset,
             {
-                'shape': (3, 32, 32),
+                'shape': (3, 224, 224),
             },
         ),
         'kwargs': {
-            'stochastic_method': 'residual',
-            'target_layer_name': 'BasicBlock',
+            'stochastic_method': 'block',
+            'target_layer_name': 'ResNetBottleneck',
             'drop_rate': 0.2,
             'drop_distribution': 'linear',
             'drop_warmup': '0.0dur',

diff --git a/tests/algorithms/test_algorithms_train.py b/tests/algorithms/test_algorithms_train.py
@@ -18,7 +18,7 @@ def test_algorithm_trains(alg_cls: type[Algorithm]):
     trainer = Trainer(
         model=model,
         train_dataloader=dataloader,
-        max_duration='1ep',
+        max_duration='1ba',
         algorithms=alg_cls(**alg_kwargs),
     )
     trainer.fit()