From 7cfcc76920a443f0dfc85a25576b240d5b7523e8 Mon Sep 17 00:00:00 2001
From: multivitamin <36110407+anvuongb@users.noreply.github.com>
Date: Tue, 29 Oct 2024 19:43:06 -0700
Subject: [PATCH] fix train.py on torch 2

---
 Denoising/train.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/Denoising/train.py b/Denoising/train.py
index b2c58cd..9f12931 100644
--- a/Denoising/train.py
+++ b/Denoising/train.py
@@ -124,8 +124,10 @@
         restored = model_restoration(input_)
 
         # Compute loss at each stage
-        loss = torch.sum([criterion(torch.clamp(restored[j],0,1),target) for j in range(len(restored))])
-         
+        loss = 0
+        for j in range(len(restored)):
+            lost = loss + criterion(torch.clamp(restored[j],0,1),target)
+      
         loss.backward()
         optimizer.step()
         epoch_loss +=loss.item()