diff --git a/lib/train/losses.py b/lib/train/losses.py
index f4f214d..eae0f77 100644
--- a/lib/train/losses.py
+++ b/lib/train/losses.py
@@ -74,7 +74,7 @@ class LossBalancer:
         output_params = torch.autograd.grad(current_loss, model_params, grad_outputs=torch.ones_like(current_loss, device=model_params.device), retain_graph=True, only_inputs=True)[0]
 
         # Compute L2 gradient norm
-        grad_norm = output_params.view(output_params.size(0),-1).norm(2, dim=-1)
+        grad_norm = output_params.view(output_params.size(0),-1).norm(2, dim=-1).mean()
         ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon
         return grad_norm/ema_loss
 
@@ -121,7 +121,7 @@ class LossBalancer:
         """
         
         if not self.initial_weights: self.initial_weights = {k: 1.0 for k in losses} # initialize weights
-        if not self.active: return sum(v*self.initial_weights[k] if k in self.initial_weights else v for k,v in losses)
+        if not self.active: return sum(v*self.initial_weights[k] if k in self.initial_weights else v for k,v in losses.items())
 
         gradients = {}
         valid_losses = {}
diff --git a/training_cli.py b/training_cli.py
index e3f1d07..82aae09 100644
--- a/training_cli.py
+++ b/training_cli.py
@@ -509,7 +509,11 @@ def train_and_evaluate(
             y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat)
             with autocast(enabled=False):
                 loss_mel = F.l1_loss(y_mel, y_hat_mel)*hps.train.c_mel if hps.train.get("c_mel",0.)>0 else 0 #default 45
-                loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask)*hps.train.c_kl if hps.train.get("c_kl",0.)>0 else 0 #default 1
+                if hps.train.get("c_kl",0.)>0:
+                    loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask)*hps.train.c_kl
+                else:
+                    loss_kl = F.mse_loss(z_p,m_p)+F.smooth_l1_loss(logs_q,logs_p)
+                    loss_kl /= torch.sum(z_mask)
                 loss_fm = feature_loss(fmap_r, fmap_g)*hps.train.c_fm if hps.train.get("c_fm",0.)>0 else 0 #default 2
                 harmonic_loss, tefs_loss, tsi_loss = combined_aux_loss(
                     wave, y_hat,n_mels=hps.data.n_mel_channels,sample_rate=hps.data.sampling_rate,