瀏覽代碼

update utils

imClumsyPanda 2 年之前
父節點
當前提交
3f43199449
共有 1 個文件被更改,包括 1 次插入3 次删除
  1. 1 3
      utils/__init__.py

+ 1 - 3
utils/__init__.py

@@ -1,5 +1,4 @@
-import torch.cuda
-import torch.backends
+import torch
 
 def torch_gc(DEVICE):
     if torch.cuda.is_available():
@@ -8,7 +7,6 @@ def torch_gc(DEVICE):
             torch.cuda.ipc_collect()
     elif torch.backends.mps.is_available():
         try:
-            import torch.mps
             torch.mps.empty_cache()
         except Exception as e:
             print(e)