diff --git a/networks/resnet_css.py b/networks/resnet_css.py
index dd07861..5015f59 100644
--- a/networks/resnet_css.py
+++ b/networks/resnet_css.py
@@ -97,6 +97,7 @@ def forward(self, x):
 
 
 def _freeze_module(module):
+    module.eval()
     for param in module.parameters():
         param.requires_grad = False
 
@@ -152,10 +153,15 @@ def __init__(self, block, layers, num_classes=2):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
+    def freeze(self):    
         # Freeze first layers
         _freeze_module(self.conv1)
         _freeze_module(self.bn1)
         _freeze_module(self.layer1)
+        
+    def train(self, mode=True):
+        super(ResNet, self).train(mode)
+        freeze()
 
     def _make_layer(self, block, planes, blocks, stride=1):
         downsample = None