Update src/loss.py

src/loss.py

@@ -116,8 +116,8 @@ class LossSchedulerModel(torch.nn.Module):
 		jitter=A.matrix_generator.forward(t)
 		for(D,E,J)in zip(B,A.we[t],jitter[t]):
 			dynamic_scale = torch.tanh(A.decay_rate * torch.tensor(t)).to(D.device)
-			C+=D*(E+J.to(D.device))
-			C += C * dynamic_scale
+			correction=D*(E+J.to(D.device))
+			C += correction * dynamic_scale
 		return C.to(xT.dtype) #C.to(xT.dtype)
 
 class LossScheduler: