Update short_context_recovery method to use the correct lambda_factor…

…s & n_hat
jshuadvd · Jul 2, 2024 · c8e557b · c8e557b
1 parent 85d374e
commit c8e557b
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/main.py b/src/main.py
@@ -708,6 +708,11 @@ def short_context_recovery(model, data, base_length, lambda_factors_base, n_hat_
         # Fine-tune for short context recovery
         model = fine_tune(model, data, length, lambda_factors, n_hat, steps=100)
 
+        # Update model attributes
+        key = "4k" if length == 4096 else "8k"
+        model.lambda_factors[key] = lambda_factors
+        model.n_hat[key] = n_hat
+
     # Store base factors for use during inference
     model.lambda_factors_base = lambda_factors_base
     model.n_hat_base = n_hat_base