[hotfix] pass test_complete_workflow (#1877)

51597f6a · Jiarui Fang · GitHub · 986f8cba · 51597f6a
Unverified Commit 51597f6a authored Nov 10, 2022 by Jiarui Fang Committed by GitHub Nov 10, 2022
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

tests/test_fx/test_complete_workflow.py tests/test_fx/test_complete_workflow.py +2 -2

No files found.
--- a/tests/test_fx/test_complete_workflow.py
+++ b/tests/test_fx/test_complete_workflow.py
@@ -50,7 +50,7 @@ def run_workflow(world_size, dev):
    annotated_gm.recompile()

    # materialization and sharding
-    ctx.lazy_init_parameters(annotated_gm)
+    ctx.lazy_init_parameters(annotated_gm, device=dev)
    for param in model.parameters():
        assert not param.is_meta

@@ -84,4 +84,4 @@ def test_complete_workflow(world_size, dev):


 if __name__ == '__main__':
-    test_complete_workflow(1)
+    test_complete_workflow(1, 'cuda')