diff --git a/app/controllers/budgets/investments_controller.rb b/app/controllers/budgets/investments_controller.rb
index 25425250f..490b4047d 100644
--- a/app/controllers/budgets/investments_controller.rb
+++ b/app/controllers/budgets/investments_controller.rb
@@ -89,7 +89,10 @@ module Budgets
       end
 
       def load_heading
-        @heading = @budget.headings.find(params[:heading_id]) if params[:heading_id].present?
+        if params[:heading_id].present?
+          @heading = @budget.headings.find(params[:heading_id])
+          @assigned_heading = @ballot.try(:heading_for_group, @heading.try(:group))
+        end
       end
 
   end