update test

2022-12-11 20:27:32 +08:00
parent 03965c69e6
commit 41e2858a7c
2 changed files with 3 additions and 5 deletions
--- a/lama_cleaner/tests/test_paint_by_example.py
+++ b/lama_cleaner/tests/test_paint_by_example.py
@@ -38,7 +38,7 @@ def assert_equal(
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
 def test_paint_by_example(strategy):
    model = ModelManager(name="paint_by_example", device=device)
-    cfg = get_config(strategy, paint_by_example_steps=30)
+    cfg = get_config(strategy, paint_by_example_steps=30 if device == 'cuda' else 1)
    assert_equal(
        model,
        cfg,