revery

mosaicml · Nov 30, 2024 · 868a71b · 868a71b
1 parent 7b69fc3
commit 868a71b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/trainer/test_fsdp_checkpoint.py b/tests/trainer/test_fsdp_checkpoint.py
@@ -111,11 +111,11 @@ def get_trainer(
         val_metrics=val_metrics,
     )
     model.module.to(model_init_device)
-    dataset = RandomClassificationDataset(shape=(num_features,), num_classes=num_classes, size=8)
+    dataset = RandomClassificationDataset(shape=(num_features,), num_classes=num_classes, size=128)
     dataloader = DataLoader(
         dataset,
         sampler=dist.get_sampler(dataset),
-        batch_size=2,
+        batch_size=8,
     )
     if optimizer == 'adam':
         optim = torch.optim.Adam(params=model.parameters())