Enable bfloat16 for sparse_w16a16 (#18)

neuralmagic · Feb 22, 2024 · 3c11f56 · 3c11f56
1 parent 7527b9c
commit 3c11f56
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/model_executor/layers/sparsity/sparse_w16a16.py b/vllm/model_executor/layers/sparsity/sparse_w16a16.py
@@ -31,7 +31,7 @@ def get_name(cls) -> str:
 
     @classmethod
     def get_supported_act_dtypes(cls) -> List[torch.dtype]:
-        return [torch.half]
+        return [torch.float16, torch.bfloat16]
 
     @classmethod
     def get_min_capability(cls) -> int: