Merge pull request #40 from HyeongminMoon/patch-1

L1aoXingyu · Feb 11, 2022 · 61db1de · 61db1de
2 parents 3c48d7e + db103d5
commit 61db1de
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/08-AutoEncoder/conv_autoencoder.py b/08-AutoEncoder/conv_autoencoder.py
@@ -64,8 +64,9 @@ def forward(self, x):
 criterion = nn.MSELoss()
 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate,
                              weight_decay=1e-5)
-total_loss = 0
+
 for epoch in range(num_epochs):
+    total_loss = 0
     for data in dataloader:
         img, _ = data
         img = Variable(img).cuda()
@@ -76,8 +77,8 @@ def forward(self, x):
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
+        total_loss += loss.data
     # ===================log========================
-    total_loss += loss.data
     print('epoch [{}/{}], loss:{:.4f}'
           .format(epoch+1, num_epochs, total_loss))
     if epoch % 10 == 0: