Browse Source
Check if cuda available
Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
main
1 changed files with
5 additions and
1 deletions
-
clip.py
|
|
@ -96,7 +96,11 @@ class Clip(NNOperator): |
|
|
def __init__(self, model_name: str, modality: str, device: str = 'cpu', checkpoint_path: str = None): |
|
|
def __init__(self, model_name: str, modality: str, device: str = 'cpu', checkpoint_path: str = None): |
|
|
self.model_name = model_name |
|
|
self.model_name = model_name |
|
|
self.modality = modality |
|
|
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 |
|
|
self.checkpoint_path = checkpoint_path |
|
|
config = self._configs() |
|
|
config = self._configs() |
|
|
real_name = model_name |
|
|
real_name = model_name |
|
|
|
