prevent 'torch' has no attribute 'cuda' error in torch_gc, so user can use CPU or WebGPU (#302)

2024-03-11 19:04:19 +08:00 · 2024-03-11 19:04:19 +08:00 · bb1a6191b0
commit bb1a6191b0
parent dd89041f72
1 changed files with 8 additions and 5 deletions
--- a/backend-python/utils/torch.py
+++ b/backend-python/utils/torch.py
@ -19,9 +19,12 @@ def set_torch():


 def torch_gc():
+    try:
        import torch

        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