Specifically import torch.cuda.comm in breakmodel.py

2021-08-20 10:47:54 -04:00 · 2021-08-20 10:47:54 -04:00 · eef0db8dee
parent b1c13f832a
commit eef0db8dee
1 changed files with 1 additions and 0 deletions
--- a/breakmodel.py
+++ b/breakmodel.py
@ -212,6 +212,7 @@ Released under the Apache License 2.0


 import torch
+import torch.cuda.comm
 import copy
 import gc