From bb1a6191b02c920c38af1c096a594bf8829a946a Mon Sep 17 00:00:00 2001 From: josc146 Date: Mon, 11 Mar 2024 19:04:19 +0800 Subject: [PATCH] prevent 'torch' has no attribute 'cuda' error in torch_gc, so user can use CPU or WebGPU (#302) --- backend-python/utils/torch.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/backend-python/utils/torch.py b/backend-python/utils/torch.py index 426d03a..637c391 100644 --- a/backend-python/utils/torch.py +++ b/backend-python/utils/torch.py @@ -19,9 +19,12 @@ def set_torch(): def torch_gc(): - import torch + try: + import torch - if torch.cuda.is_available(): - with torch.cuda.device(0): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() + if torch.cuda.is_available(): + with torch.cuda.device(0): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + except: + pass # prevent 'torch' has no attribute 'cuda' error, so user can use CPU or WebGPU