diff --git a/ml/src/trainers/tft.rs b/ml/src/trainers/tft.rs
index 1234567..89abcde 100644
--- a/ml/src/trainers/tft.rs
+++ b/ml/src/trainers/tft.rs
@@ -304,8 +304,8 @@ impl TFTTrainer {
         // Initialize model
         let model = TemporalFusionTransformer::new(model_config.clone())?;
 
-        // Create variable map for model parameters
-        let var_map = Arc::new(VarMap::new());
+        // Get variable map from model (contains all model weights)
+        let var_map = model.get_varmap().clone();
 
         // Create checkpoint manager with proper CheckpointConfig
         let checkpoint_config = CheckpointConfig {