Update.

unsuper/minist.py

@@ -125,14 +125,14 @@ for epoch in range(epochs):
         loss.backward()
         optimizer_unsuper.step()
         if (i + 1) % 100 == 0:
-            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.8f}")
+            print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.8f}")
 
 # Train the model
 model.conv1.weight.requires_grad = False
 model.conv2.weight.requires_grad = True
 model.fc1.weight.requires_grad = True
 criterion = nn.CrossEntropyLoss()
-optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.6)
+optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.2)
 n_total_steps = len(train_loader)
 for epoch in range(num_epochs):
     for i, (images, labels) in enumerate(train_loader):