ugh

2024-12-06 22:35:30 -06:00 · 2024-12-06 22:35:30 -06:00 · 953d3eb030
commit 953d3eb030
parent 42fafbaaca
1 changed files with 1 additions and 1 deletions
--- a/vall_e/models/base.py
+++ b/vall_e/models/base.py
@ -1671,9 +1671,9 @@ class Base(nn.Module):

 				# mix if not nan
 				if not torch.isnan(soft_loss).any():
-					loss['kl'] = soft_loss * A
 					for k in loss.keys():
 						loss[k] *= (1.0 - A)
+					loss['kl'] = soft_loss * A

 			# include any additional losses (for example: MoE router)
 			if output.loss is not None: