still hadn't fixed the issue, now fixed

jbloomAus · Feb 1, 2024 · a36ee21 · a36ee21
1 parent b4546db
commit a36ee21
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sae_training/sparse_autoencoder.py b/sae_training/sparse_autoencoder.py
@@ -125,10 +125,11 @@ def forward(self, x, dead_neuron_mask = None):
             mse_rescaling_factor = (mse_loss / mse_loss_ghost_resid).detach()
             mse_loss_ghost_resid = mse_rescaling_factor * mse_loss_ghost_resid
 
+        mse_loss_ghost_resid = mse_loss_ghost_resid.mean()
         mse_loss = mse_loss.mean()
         sparsity = torch.abs(feature_acts).sum(dim=1).mean(dim=(0,)) 
         l1_loss = self.l1_coefficient * sparsity
-        loss = mse_loss + l1_loss + mse_loss_ghost_resid.mean()
+        loss = mse_loss + l1_loss + mse_loss_ghost_resid
 
         return sae_out, feature_acts, loss, mse_loss, l1_loss, mse_loss_ghost_resid