[zero] fix testing parameters (#2042)

384cd263 · HELSON · GitHub · 17a3c685 · 384cd263
Unverified Commit 384cd263 authored Nov 30, 2022 by HELSON Committed by GitHub Nov 30, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/test_gemini/update/test_optim.py tests/test_gemini/update/test_optim.py +1 -1

No files found.
--- a/tests/test_gemini/update/test_optim.py
+++ b/tests/test_gemini/update/test_optim.py
@@ -44,7 +44,7 @@ TEST_MODELS = ['gpt2', 'bert']
 EXAMPLE_MODELS = ['simple_net']


-@parameterize('placement_policy', ['cuda'])
+@parameterize('placement_policy', ['cuda', 'cpu', 'auto', 'const'])
 @parameterize('model_name', TEST_MODELS)
 def exam_model_step(placement_policy, model_name: str):
    set_seed(42)