From 527aff4f6975b353f47b3225bbc16ea41b4a81df Mon Sep 17 00:00:00 2001 From: "junjie.jiang" Date: Mon, 15 Apr 2024 14:36:45 +0800 Subject: [PATCH] Check gpu is available Signed-off-by: junjie.jiang --- timm_image.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/timm_image.py b/timm_image.py index 026625f..e7f936e 100644 --- a/timm_image.py +++ b/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(