diff --git a/.idea/mmsr.iml b/.idea/mmsr.iml
index 2d68ad0f..dc3935a0 100644
--- a/.idea/mmsr.iml
+++ b/.idea/mmsr.iml
@@ -11,7 +11,6 @@
     </content>
     <orderEntry type="jdk" jdkName="torch-cu11" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
-    <orderEntry type="module" module-name="SwitchedConvolutions" />
   </component>
   <component name="PyDocumentationSettings">
     <option name="format" value="PLAIN" />
diff --git a/.idea/modules.xml b/.idea/modules.xml
index fdaa2758..e88fd2ba 100644
--- a/.idea/modules.xml
+++ b/.idea/modules.xml
@@ -2,7 +2,6 @@
 <project version="4">
   <component name="ProjectModuleManager">
     <modules>
-      <module fileurl="file://$PROJECT_DIR$/../SwitchedConvolutions/.idea/SwitchedConvolutions.iml" filepath="$PROJECT_DIR$/../SwitchedConvolutions/.idea/SwitchedConvolutions.iml" />
       <module fileurl="file://$PROJECT_DIR$/.idea/mmsr.iml" filepath="$PROJECT_DIR$/.idea/mmsr.iml" />
     </modules>
   </component>
diff --git a/codes/models/archs/StructuredSwitchedGenerator.py b/codes/models/archs/StructuredSwitchedGenerator.py
index 6533719a..b5b37500 100644
--- a/codes/models/archs/StructuredSwitchedGenerator.py
+++ b/codes/models/archs/StructuredSwitchedGenerator.py
@@ -10,7 +10,7 @@ from switched_conv.switched_conv_util import save_attention_to_image_rgb
 from switched_conv.switched_conv import compute_attention_specificity
 import os
 import torchvision
-from torch.utils.checkpoint import checkpoint
+from utils.util import checkpoint
 
 # VGG-style layer with Conv(stride2)->BN->Activation->Conv->BN->Activation
 # Doubles the input filter count.