From c0f26b5975fac1a10951f4d0d893a28ff10c7f78 Mon Sep 17 00:00:00 2001
From: "junjie.jiang" <junjie.jiang@zilliz.com>
Date: Mon, 15 Apr 2024 14:18:59 +0800
Subject: [PATCH] Check if cuda available

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
---
 clip.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/clip.py b/clip.py
index 8bedb00..cf110dc 100644
--- a/clip.py
+++ b/clip.py
@@ -96,7 +96,11 @@ class Clip(NNOperator):
     def __init__(self, model_name: str, modality: str, device: str = 'cpu', checkpoint_path: str = None):
         self.model_name = model_name
         self.modality = modality
-        self.device = device
+        if not torch.cuda.is_available():
+            log.warning("Gpu not available, use cpu")
+            self.device = 'cpu'
+        else:
+            self.device = device
         self.checkpoint_path = checkpoint_path
         config = self._configs()
         real_name = model_name