diff --git a/unsuper/conv1_output.png b/unsuper/conv1_output.png
index 69eda5f..338397c 100644
Binary files a/unsuper/conv1_output.png and b/unsuper/conv1_output.png differ
diff --git a/unsuper/conv1_weight.png b/unsuper/conv1_weight.png
index f8e427a..76c8669 100644
Binary files a/unsuper/conv1_weight.png and b/unsuper/conv1_weight.png differ
diff --git a/unsuper/conv1_weight_grad.png b/unsuper/conv1_weight_grad.png
index 9b14c01..ab56018 100644
Binary files a/unsuper/conv1_weight_grad.png and b/unsuper/conv1_weight_grad.png differ
diff --git a/unsuper/conv2_output.png b/unsuper/conv2_output.png
index d83e49b..db24d41 100644
Binary files a/unsuper/conv2_output.png and b/unsuper/conv2_output.png differ
diff --git a/unsuper/conv2_weight.png b/unsuper/conv2_weight.png
index 3bba953..82eab24 100644
Binary files a/unsuper/conv2_weight.png and b/unsuper/conv2_weight.png differ
diff --git a/unsuper/conv2_weight_grad.png b/unsuper/conv2_weight_grad.png
index 0f791a5..1840384 100644
Binary files a/unsuper/conv2_weight_grad.png and b/unsuper/conv2_weight_grad.png differ
diff --git a/unsuper/fc_output.png b/unsuper/fc_output.png
index 8d4e4d3..d6b0488 100644
Binary files a/unsuper/fc_output.png and b/unsuper/fc_output.png differ
diff --git a/unsuper/fc_weight.png b/unsuper/fc_weight.png
index 65e00b5..2164913 100644
Binary files a/unsuper/fc_weight.png and b/unsuper/fc_weight.png differ
diff --git a/unsuper/fc_weight_grad.png b/unsuper/fc_weight_grad.png
index ed58f19..aa41d81 100644
Binary files a/unsuper/fc_weight_grad.png and b/unsuper/fc_weight_grad.png differ
diff --git a/unsuper/minist.py b/unsuper/minist.py
index ae3834c..d3b68ad 100644
--- a/unsuper/minist.py
+++ b/unsuper/minist.py
@@ -13,8 +13,8 @@ seed = 4321
 torch.manual_seed(seed)
 torch.cuda.manual_seed_all(seed)
 
-# Device configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("mps")
 
 # Hyper-parameters
 num_epochs = 1
@@ -51,25 +51,25 @@ class ConvNet(nn.Module):
         return x
 
     def printFector(self, x, label):
-        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "input_image.png")
+        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "input_image.png")
         x = self.conv1(x)
         w = self.conv1.weight
-        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight.png")
+        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight.png")
 
-        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "conv1_output.png")
+        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv1_output.png")
         x = self.pool(F.relu(x))
         x = self.conv2(x)
         w = self.conv2.weight
-        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv2_weight.png")
+        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv2_weight.png")
 
-        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), "conv2_output.png")
+        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), "conv2_output.png")
         x = self.pool(F.relu(x))
         x = x.view(-1, 8 * 5 * 5)
         x = self.fc1(x)
 
-        show.DumpTensorToImage(self.fc1.weight.view(-1, 10, 10).permute(2, 0, 1), "fc_weight.png")
+        show.DumpTensorToImage(self.fc1.weight.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight.png")
 
-        show.DumpTensorToImage(x.view(-1), "fc_output.png")
+        show.DumpTensorToImage(x.view(-1).cpu(), "fc_output.png")
 
         criterion = nn.CrossEntropyLoss()
         loss = criterion(x, label)
@@ -77,10 +77,10 @@ class ConvNet(nn.Module):
         loss.backward()
 
         w = self.conv1.weight.grad
-        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight_grad.png")
+        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight_grad.png")
         w = self.conv2.weight.grad
-        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv2_weight_grad.png")
-        show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 10, 10).permute(2, 0, 1), "fc_weight_grad.png")
+        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv2_weight_grad.png")
+        show.DumpTensorToImage(self.fc1.weight.grad.view(-1, 10, 10).permute(2, 0, 1).cpu(), "fc_weight_grad.png")
 
 
 model = ConvNet().to(device)
@@ -109,6 +109,8 @@ for epoch in range(num_epochs):
             print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}")
 
 for images, labels in test_loader:
+    images = images.to(device)
+    labels = labels.to(device)
     model.printFector(images, labels)
     break