From eef0db8deec1735c5b8eabe6e786596101b6a1b3 Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Fri, 20 Aug 2021 10:47:54 -0400
Subject: [PATCH] Specifically import torch.cuda.comm in breakmodel.py

---
 breakmodel.py | 1 +
 1 file changed, 1 insertion(+)

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