Fix loss cloning

ai-safety-foundation · Jan 30, 2024 · 06edb9e · 06edb9e
1 parent e30dd43
commit 06edb9e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sparse_autoencoder/train/pipeline.py b/sparse_autoencoder/train/pipeline.py
@@ -498,7 +498,7 @@ def run_pipeline(
         )
 
         # Get the loss fn
-        loss_fn = self.autoencoder.loss_metric.clone()
+        loss_fn = self.loss.clone()
         loss_fn.keep_batch_dim = True
 
         with tqdm(