diff --git a/lib/mix/tasks/propagation_train.ex b/lib/mix/tasks/propagation_train.ex
index e33ab2a0..1e9eba4c 100644
--- a/lib/mix/tasks/propagation_train.ex
+++ b/lib/mix/tasks/propagation_train.ex
@@ -48,7 +48,8 @@ defmodule Mix.Tasks.PropagationTrain do
           pretrain_epochs: :integer,
           finetune_epochs: :integer,
           pretrain_sample: :integer,
-          batch_size: :integer
+          batch_size: :integer,
+          no_pretrain: :boolean
         ]
       )
 
@@ -68,38 +69,53 @@ defmodule Mix.Tasks.PropagationTrain do
     finetune_epochs = Keyword.get(opts, :finetune_epochs, 200)
     pretrain_sample = Keyword.get(opts, :pretrain_sample, 500_000)
     batch_size = Keyword.get(opts, :batch_size, 256)
+    skip_pretrain = Keyword.get(opts, :no_pretrain, false)
 
     "=" |> String.duplicate(70) |> IO.puts()
-    IO.puts("PROPAGATION MODEL TRAINING (Two-Phase, 20 features)")
+    IO.puts("PROPAGATION MODEL TRAINING (20 features)")
     "=" |> String.duplicate(70) |> IO.puts()
-    IO.puts("Phase 1: #{pretrain_epochs} epochs on #{pretrain_sample} algorithm scores")
-    IO.puts("Phase 2: #{finetune_epochs} epochs on real QSO data")
-    IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")
 
-    # ── Phase 1: Pre-train on algorithm scores ──────────────────────
-    IO.puts("=" <> String.duplicate("─", 50))
-    IO.puts("PHASE 1: Pre-training on algorithm scores + solar + soundings")
-    IO.puts("=" <> String.duplicate("─", 50))
+    pretrained_state =
+      if skip_pretrain do
+        IO.puts("Skipping pre-training (--no-pretrain)")
+        nil
+      else
+        IO.puts("Phase 1: #{pretrain_epochs} epochs on #{pretrain_sample} algorithm scores")
+        IO.puts("Phase 2: #{finetune_epochs} epochs on real QSO data")
+        IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")
 
-    {pt_features, pt_targets, pt_counts} = load_pretrain_data(pretrain_sample)
-    pt_n = elem(Nx.shape(pt_features), 0)
-    IO.puts("Loaded #{pt_n} samples:")
-    print_band_counts(pt_counts)
+        # ── Phase 1: Pre-train on algorithm scores ──────────────────
+        IO.puts("=" <> String.duplicate("─", 50))
+        IO.puts("PHASE 1: Pre-training on algorithm scores + solar + soundings")
+        IO.puts("=" <> String.duplicate("─", 50))
 
-    {pt_features, pt_targets} = shuffle(pt_features, pt_targets, pt_n)
-    {pt_train_x, pt_train_y, pt_val_x, pt_val_y, _, _} = split(pt_features, pt_targets, pt_n)
+        {pt_features, pt_targets, pt_counts} = load_pretrain_data(pretrain_sample)
+        pt_n = elem(Nx.shape(pt_features), 0)
+        IO.puts("Loaded #{pt_n} samples:")
+        print_band_counts(pt_counts)
 
-    IO.puts("\nPre-training for #{pretrain_epochs} epochs (lr=0.001)...")
+        {pt_features, pt_targets} = shuffle(pt_features, pt_targets, pt_n)
+        {pt_train_x, pt_train_y, pt_val_x, pt_val_y, _, _} = split(pt_features, pt_targets, pt_n)
 
-    {pretrained_state, pt_metrics} =
-      Model.train(pt_train_x, pt_train_y, epochs: pretrain_epochs, batch_size: batch_size)
+        IO.puts("\nPre-training for #{pretrain_epochs} epochs (lr=0.001)...")
 
-    print_loss(pt_metrics.final_loss)
-    print_eval("Pre-train val", Model.evaluate(pretrained_state, pt_val_x, pt_val_y))
+        {state, pt_metrics} =
+          Model.train(pt_train_x, pt_train_y, epochs: pretrain_epochs, batch_size: batch_size)
 
-    # ── Phase 2: Fine-tune on real QSO data ─────────────────────────
+        print_loss(pt_metrics.final_loss)
+        print_eval("Pre-train val", Model.evaluate(state, pt_val_x, pt_val_y))
+        state
+      end
+
+    # ── QSO training ────────────────────────────────────────────────
     IO.puts("\n" <> "=" <> String.duplicate("─", 50))
-    IO.puts("PHASE 2: Fine-tuning on QSO + HRRR + solar + soundings")
+
+    if pretrained_state do
+      IO.puts("PHASE 2: Fine-tuning on QSO + HRRR + solar + soundings")
+    else
+      IO.puts("TRAINING: QSO + HRRR + solar + soundings (direct, lr=0.001)")
+    end
+
     IO.puts("=" <> String.duplicate("─", 50))
 
     {ft_features, ft_targets, ft_counts} = load_qso_data()
@@ -110,15 +126,13 @@ defmodule Mix.Tasks.PropagationTrain do
     {ft_features, ft_targets} = shuffle(ft_features, ft_targets, ft_n)
     {ft_train_x, ft_train_y, ft_val_x, ft_val_y, ft_test_x, ft_test_y} = split(ft_features, ft_targets, ft_n)
 
-    IO.puts("\nFine-tuning for #{finetune_epochs} epochs (lr=0.0003)...")
+    lr = if pretrained_state, do: 0.0003, else: 0.001
+    IO.puts("\nTraining for #{finetune_epochs} epochs (lr=#{lr})...")
 
-    {trained_state, ft_metrics} =
-      Model.train(ft_train_x, ft_train_y,
-        epochs: finetune_epochs,
-        batch_size: batch_size,
-        learning_rate: 0.0003,
-        initial_state: pretrained_state
-      )
+    train_opts = [epochs: finetune_epochs, batch_size: batch_size, learning_rate: lr]
+    train_opts = if pretrained_state, do: [{:initial_state, pretrained_state} | train_opts], else: train_opts
+
+    {trained_state, ft_metrics} = Model.train(ft_train_x, ft_train_y, train_opts)
 
     print_loss(ft_metrics.final_loss)
     print_eval("Fine-tune val", Model.evaluate(trained_state, ft_val_x, ft_val_y))