From fe4ef95d928222ab66f886b70c0ee508408578f4 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Mon, 6 May 2024 20:32:08 +0200
Subject: [PATCH] Protect cuda better ?

---
 server/text_generation_server/utils/import_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/text_generation_server/utils/import_utils.py b/server/text_generation_server/utils/import_utils.py
index 2cabfe36..f54987eb 100644
--- a/server/text_generation_server/utils/import_utils.py
+++ b/server/text_generation_server/utils/import_utils.py
@@ -29,7 +29,7 @@ if torch.version.hip is not None:
     empty_cache = torch.cuda.empty_cache
     synchronize = torch.cuda.synchronize
     get_free_memory = get_cuda_free_memory
-elif torch.version.cuda is not None:
+elif torch.version.cuda is not None and torch.cuda.is_available():
     SYSTEM = "cuda"
     empty_cache = torch.cuda.empty_cache
     synchronize = torch.cuda.synchronize