diff --git a/unsuper/minist.py b/unsuper/minist.py
index 52bb775..5c984ac 100644
--- a/unsuper/minist.py
+++ b/unsuper/minist.py
@@ -36,22 +36,23 @@ class ConvNet(nn.Module):
         self.fc1 = nn.Linear(1 * 4 * 4, 10)
 
     def forward(self, x):
-        x = self.pool(self.conv1(x))
+        x = self.forward_unsuper(x)
+        x = self.pool(x)
         x = self.pool(self.conv2(x))
         x = x.view(x.shape[0], -1)
         x = self.fc1(x)
         return x
 
     def forward_unsuper(self, x):
-        x = self.conv1(x)
-        # x = self.pool(self.conv1(x))
-        return x
-
-    def forward_finetune(self, x):
-        x = self.pool(self.conv1(x))
-        x = self.pool(self.conv2(x))
-        x = x.view(x.shape[0], -1)
-        x = self.fc1(x)
+        weight = self.conv1.weight.reshape(self.conv1.weight.shape[0], -1)
+        weight = weight.permute(1, 0)
+        mean = torch.mean(weight, dim=0)
+        weight = weight - mean
+        sum = torch.sum(torch.abs(weight), dim=0)
+        weight = weight / sum
+        weight = weight.permute(1, 0)
+        weight = weight.reshape(self.conv1.weight.shape)
+        x = torch.conv2d(x, weight)
         return x
 
     def printFector(self, x, label, dir=""):
@@ -104,7 +105,7 @@ model = ConvNet().to(device)
 model.train()
 
 # Train the model unsuper
-epochs = 2
+epochs = 10
 model.conv1.weight.requires_grad = True
 model.conv2.weight.requires_grad = False
 model.fc1.weight.requires_grad = False
@@ -114,24 +115,40 @@ for epoch in range(epochs):
         images = images.to(device)
         outputs = model.forward_unsuper(images)
 
-        outputs = outputs.permute(1, 0, 2, 3)  # 64 8 24 24 -> 8 64 24 24
-        sample = outputs.reshape(outputs.shape[0], -1)  # -> 8 36864
+        # outputs = outputs.permute(0, 2, 3, 1)  # 64 8 24 24 -> 64 24 24 8
+        # sample = outputs.reshape(-1, outputs.shape[3])  # -> 36864 8
+        # abs = torch.abs(sample)
+        # max, max_index = torch.max(abs, dim=1)
+        # min, min_index = torch.min(abs, dim=1)
+        # label = sample * 0.9
+        # all = range(0, label.shape[0])
+        # label[all, max_index] = label[all, max_index]*1.1
+        # loss = F.l1_loss(sample, label)
+        # model.conv1.weight.grad = None
+        # loss.backward()
 
-        sample_mean = torch.mean(sample, dim=1, keepdim=True)
-        diff_mean = torch.mean(torch.abs(sample - sample_mean), dim=1, keepdim=True)
-        diff_ratio = (sample - sample_mean) / diff_mean
-        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)
+        outputs = outputs.permute(0, 2, 3, 1)  # 64 8 24 24 -> 64 24 24 8
+        sample = outputs.reshape(outputs.shape[0], -1, outputs.shape[3])  # -> 64 24x24 8
+        abs = torch.abs(sample)
+        sum = torch.sum(abs, dim=1, keepdim=False)
+        max, max_index = torch.max(sum, dim=1)
+        label = sample * 0.9
+        all = range(0, label.shape[0])
+        all_wh = range(0, 24 * 24)
+        label[all, :, max_index] = label[all, :, max_index] * 1.1
+        loss = F.l1_loss(sample, label)
         model.conv1.weight.grad = None
         loss.backward()
 
-        grad = model.conv1.weight.data
-        grad = grad.view(8, -1)
-        grad_mean = torch.mean(grad, dim=1)
-        max, index = torch.max(grad_mean, dim=0)
+        # show.DumpTensorToImage(images.view(-1, images.shape[2], images.shape[3]), "input_image.png", Contrast=[0, 1.0])
+        # w = model.conv1.weight.data
+        # show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight.png", Contrast=[-1.0, 1.0])
+        # w = model.conv1.weight.grad
+        # show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight_grad.png")
+        model.conv1.weight.data = model.conv1.weight.data - model.conv1.weight.grad * 1000
+        # w = model.conv1.weight.data
+        # show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight_update.png", Contrast=[-1.0, 1.0])
 
-        model.conv1.weight.data[index] = model.conv1.weight.data[index] - model.conv1.weight.grad[index] * 0.2
         if (i + 1) % 100 == 0:
             print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.8f}")
 
@@ -146,7 +163,7 @@ for epoch in range(num_epochs):
     for i, (images, labels) in enumerate(train_loader):
         images = images.to(device)
         labels = labels.to(device)
-        outputs = model.forward_finetune(images)
+        outputs = model(images)
         loss = criterion(outputs, labels)
         optimizer.zero_grad()
         loss.backward()
@@ -154,7 +171,7 @@ for epoch in range(num_epochs):
         if (i + 1) % 100 == 0:
             print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}")
 
-print("Finished Training")
+# print("Finished Training")
 
 test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
 test_loader = iter(test_loader)