Update como.py

zhenye234 · Oct 23, 2023 · 255856e · 255856e
1 parent 01efbe3
commit 255856e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/model/como.py b/model/como.py
@@ -269,7 +269,7 @@ def EDMLoss(self, x_start,   cond,nonpadding ):
         n = (torch.randn_like(x_start)+cond ) * sigma
         D_yn = self.EDMPrecond(x_start + n, sigma ,cond,self.denoise_fn,nonpadding)
         loss = (weight * ((D_yn - x_start) ** 2))
-        loss=loss*nonpadding.unsqueeze(1).unsqueeze(1)
+        loss=loss*nonpadding 
         loss=loss.mean() 
         return loss
 
@@ -361,7 +361,7 @@ def CTLoss_D(self,y, cond,nonpadding): #k
 
 
         loss =   (f_theta - f_theta_ema.detach()) ** 2 
-        loss=loss*nonpadding.unsqueeze(1).unsqueeze(1)
+        loss=loss*nonpadding 
         loss=loss.mean()