avoid division by 0

jbloomAus · Apr 10, 2024 · 2c7c6d8 · 2c7c6d8
1 parent e1ad1aa
commit 2c7c6d8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sae_lens/training/sparse_autoencoder.py b/sae_lens/training/sparse_autoencoder.py
@@ -377,4 +377,6 @@ def _per_item_mse_loss_with_target_norm(
     """
     target_centered = target - target.mean(dim=0, keepdim=True)
     normalization = target_centered.norm(dim=-1, keepdim=True)
-    return torch.nn.functional.mse_loss(preds, target, reduction="none") / normalization
+    return torch.nn.functional.mse_loss(preds, target, reduction="none") / (
+        normalization + 1e-6
+    )