Browse Source
Check gpu is available
Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
main
1 changed files with
8 additions and
3 deletions
-
timm_image.py
|
|
@ -89,8 +89,13 @@ class TimmImage(NNOperator): |
|
|
checkpoint_path: str = None |
|
|
checkpoint_path: str = None |
|
|
) -> None: |
|
|
) -> None: |
|
|
super().__init__() |
|
|
super().__init__() |
|
|
|
|
|
if not torch.cuda.is_available(): |
|
|
|
|
|
log.warning('Gpu is not available, use cpu') |
|
|
|
|
|
self.device = 'cpu' |
|
|
|
|
|
else: |
|
|
if device is None: |
|
|
if device is None: |
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
|
|
|
|
|
self.device = 'cuda' |
|
|
|
|
|
else: |
|
|
self.device = device |
|
|
self.device = device |
|
|
self.model_name = model_name |
|
|
self.model_name = model_name |
|
|
if self.model_name: |
|
|
if self.model_name: |
|
|
|
