diff --git a/breakmodel.py b/breakmodel.py
index 163681ad..7417db36 100644
--- 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