prevent 'torch' has no attribute 'cuda' error in torch_gc, so user ca…

…n use CPU or WebGPU (#302)
josStorer · Mar 11, 2024 · bb1a619 · bb1a619
1 parent dd89041
commit bb1a619
Showing 1 changed file with 8 additions and 5 deletions.
diff --git 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