Update tests/callbacks/test_oom_observer.py

Co-authored-by: Mihir Patel <[email protected]>
mosaicml · Feb 2, 2024 · 78bce44 · 78bce44
1 parent 7b7f30c
commit 78bce44
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/tests/callbacks/test_oom_observer.py b/tests/callbacks/test_oom_observer.py
@@ -17,15 +17,11 @@
 
 @pytest.mark.skipif(version.parse(torch.__version__) < version.parse('2.1.0'),
                     reason='OOM Observer requires PyTorch 2.1 or higher')
-def test_oom_observer_warnings_on_cpu_models(device: str):
-
-    # Error if the user sets device=cpu even when cuda is available
-    del device  # unused. always using cpu
+def test_oom_observer_warnings_on_cpu_models():
     ob = OOMObserver()
     Trainer(
         model=SimpleModel(),
         callbacks=ob,
-        device='cpu',
         train_dataloader=DataLoader(RandomClassificationDataset()),
         max_duration='1ba',
     )