Applied the format changes in conftest

HabanaAI · Sep 30, 2024 · b0a9d02 · b0a9d02
1 parent f59495a
commit b0a9d02
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py
@@ -84,10 +84,7 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool):
 @pytest.fixture
 def dist_init():
     temp_file = tempfile.mkstemp()[1]
-    if is_hpu():
-        backend_type = "hccl"
-    else:
-        backend_type = "nccl"
+    backend_type = "hccl" if is_hpu() else "nccl"
     init_distributed_environment(
         world_size=1,
         rank=0,
@@ -262,8 +259,10 @@ def get_model_patched(*, model_config, device_config, **kwargs):
         return get_model_old(model_config=model_config,
                              device_config=device_config,
                              **kwargs)
+
     if is_hpu():
-        with patch("vllm.worker.habana_model_runner.get_model", get_model_patched):
+        with patch("vllm.worker.habana_model_runner.get_model",
+                   get_model_patched):
             engine = vllm.LLM("meta-llama/Llama-2-7b-hf", enable_lora=False)
     else:
         with patch("vllm.worker.model_runner.get_model", get_model_patched):