diff --git a/+nosnoc/+discrete_time_problem/Gcs.m b/+nosnoc/+discrete_time_problem/Gcs.m
index 9d8dac4c..0cae58d0 100644
--- a/+nosnoc/+discrete_time_problem/Gcs.m
+++ b/+nosnoc/+discrete_time_problem/Gcs.m
@@ -132,12 +132,7 @@ function generate_direct_transcription_constraints(obj)
             
             x_prev = obj.w.x(0,0,opts.n_s);
             for ii=1:opts.N_stages
-                if obj.opts.use_fesd
-                    t_stage = obj.p.T()/opts.N_stages;
-                    h0 = obj.p.T().val/(opts.N_stages*opts.N_finite_elements(ii));
-                else
-                    h0 = obj.p.T().val/(opts.N_stages*opts.N_finite_elements(ii));
-                end
+                h0 = obj.p.T().val/(opts.N_stages*opts.N_finite_elements(ii));
                 
                 ui = obj.w.u(ii);
                 p_stage = obj.p.p_time_var(ii);
@@ -149,6 +144,13 @@ function generate_direct_transcription_constraints(obj)
                 else
                     s_sot = 1;
                 end
+                if opts.time_optimal_problem && ~opts.use_speed_of_time_variables
+                    t_stage = obj.w.T_final()/(opts.N_stages*opts.N_finite_elements(ii));
+                elseif opts.time_optimal_problem
+                    t_stage = s_sot*obj.p.T()/opts.N_stages;
+                else
+                    t_stage = obj.p.T()/opts.N_stages;
+                end
 
                 sum_h = 0;
                 for jj=1:opts.N_finite_elements(ii)