fix error in python3.8. (#2646)

InternLM · Oct 24, 2024 · d00e470 · d00e470
1 parent eaa4e6f
commit d00e470
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/lmdeploy/pytorch/engine/model_agent.py b/lmdeploy/pytorch/engine/model_agent.py
@@ -483,9 +483,8 @@ def _start_tp_process(proc_id: int,
                                 timeout=timedelta(days=35600))
         dist_ctx = DistContext(rank=rank, world_size=world_size)
         torch.cuda.set_device(rank)
-        with (get_dist_manager().context(dist_ctx),
-              get_device_manager().context(device_context),
-              torch.inference_mode()):
+        with get_dist_manager().context(dist_ctx), get_device_manager(
+        ).context(device_context), torch.inference_mode():
             args = args or tuple()
             kwargs = kwargs or dict()
             func(rank, *args, **kwargs)