p

Signed-off-by: kevin <[email protected]>
vllm-project · Sep 28, 2024 · bfee096 · bfee096
1 parent 0ad2950
commit bfee096
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/pipeline_generator/utils.py b/scripts/pipeline_generator/utils.py
@@ -32,7 +32,7 @@ class AgentQueue(str, enum.Enum):
 def get_agent_queue(no_gpu: Optional[bool], gpu_type: Optional[str], num_gpus: Optional[int]) -> AgentQueue:
     if no_gpu:
         return AgentQueue.AWS_SMALL_CPU
-    if gpu_type == A100_GPU:
+    if gpu_type == GPUType.A100.value:
         return AgentQueue.A100
     return AgentQueue.AWS_1xL4 if num_gpus == 1 else AgentQueue.AWS_4xL4