Enable bfloat16 for sparse_w16a16 (#18)

neuralmagic · Feb 21, 2024 · 6ece5c5 · 6ece5c5
1 parent c0ded57
commit 6ece5c5
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: