[Hardware][Neuron] Refactor neuron support (#3471)

2024-03-21 18:22:17 -07:00
parent ea5f14e6ff
commit e90fc21f2e
33 changed files with 615 additions and 549 deletions
--- a/tests/worker/test_swap.py
+++ b/tests/worker/test_swap.py
@ -30,7 +30,7 @@ def test_swap() -> None:
    )

    # Initialize the worker.
-    worker.init_model()
+    worker.init_device()
    worker.load_model()
    worker.init_cache_engine(cache_config)
    worker.warm_up_model()