Don't enable custom kernels if CUDA is not available (#6)
This commit is contained in:
parent
8a5f80bb61
commit
b5fadc4c28
|
@ -40,7 +40,7 @@ from text_generation_server.utils.layers import (
|
||||||
)
|
)
|
||||||
|
|
||||||
CUSTOM_KERNELS_ENABLED = False
|
CUSTOM_KERNELS_ENABLED = False
|
||||||
if not os.environ.get("DISABLE_CUSTOM_KERNELS", "False") == "True":
|
if torch.cuda.is_available() and not os.environ.get("DISABLE_CUSTOM_KERNELS", "False") == "True":
|
||||||
try:
|
try:
|
||||||
from custom_kernels import fused_bloom_attention_cuda
|
from custom_kernels import fused_bloom_attention_cuda
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue