diff --git a/unsuper/minist.py b/unsuper/minist.py
index 21abca5..ec02507 100644
--- a/unsuper/minist.py
+++ b/unsuper/minist.py
@@ -43,7 +43,8 @@ class ConvNet(nn.Module):
         return x
 
     def forward_unsuper(self, x):
-        x = self.pool(self.conv1(x))
+        x = self.conv1(x)
+        # x = self.pool(self.conv1(x))
         return x
 
     def forward_finetune(self, x):
@@ -120,9 +121,7 @@ for epoch in range(epochs):
         diff_ratio_mean = torch.mean(diff_ratio * diff_ratio, dim=1)
         label = diff_ratio_mean * 0.5
         loss = F.l1_loss(diff_ratio_mean, label)
-        if model.conv1.weight.grad is None:
-            model.conv1.weight.grad = model.conv1.weight.data
-        model.conv1.weight.grad = model.conv1.weight.grad * 0.0
+        model.conv1.weight.grad = None
         loss.backward()
         model.conv1.weight.data = model.conv1.weight.data - model.conv1.weight.grad * 0.2
         if (i + 1) % 100 == 0: