From cea79ec9c017de50d07b291d583f60eab0123eea Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 29 Mar 2024 04:00:49 +0000 Subject: [PATCH] [Misc] Minor type annotation fix --- vllm/attention/selector.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/attention/selector.py b/vllm/attention/selector.py index 42f4284c6c77..c2ec4376c9f3 100644 --- a/vllm/attention/selector.py +++ b/vllm/attention/selector.py @@ -1,4 +1,5 @@ from functools import lru_cache +from typing import Type import torch @@ -10,7 +11,7 @@ @lru_cache(maxsize=None) -def get_attn_backend(dtype: torch.dtype) -> AttentionBackend: +def get_attn_backend(dtype: torch.dtype) -> Type[AttentionBackend]: if _can_use_flash_attn(dtype): logger.info("Using FlashAttention backend.") from vllm.attention.backends.flash_attn import ( # noqa: F401