logo
Browse Source

Check gpu is available

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
main
junjie.jiang 7 months ago
parent
commit
527aff4f69
  1. 11
      timm_image.py

11
timm_image.py

@ -89,9 +89,14 @@ class TimmImage(NNOperator):
checkpoint_path: str = None
) -> None:
super().__init__()
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.device = device
if not torch.cuda.is_available():
log.warning('Gpu is not available, use cpu')
self.device = 'cpu'
else:
if device is None:
self.device = 'cuda'
else:
self.device = device
self.model_name = model_name
if self.model_name:
self.model = Model(

Loading…
Cancel
Save