diff --git a/scratch/etienne/intersimple/gail_flat.py b/scratch/etienne/intersimple/gail_flat.py
index 3996649..9714687 100644
--- a/scratch/etienne/intersimple/gail_flat.py
+++ b/scratch/etienne/intersimple/gail_flat.py
@@ -47,6 +47,7 @@ gail_trainer = adversarial.GAIL(
     n_disc_updates_per_round=32,
     discrim_kwargs={'discrim_net': MlpDiscriminator()},
     gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4530),
+    allow_variable_horizon=True,
 )
 gail_trainer.train(total_timesteps=400000)
 gail_trainer.gen_algo.save(model_name)
diff --git a/scratch/etienne/intersimple/gail_image.py b/scratch/etienne/intersimple/gail_image.py
index a1ea3d5..da6dced 100644
--- a/scratch/etienne/intersimple/gail_image.py
+++ b/scratch/etienne/intersimple/gail_image.py
@@ -47,6 +47,7 @@ gail_trainer = adversarial.GAIL(
     #n_disc_updates_per_round=2048,
     discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
     gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
+    allow_variable_horizon=True,
 )
 gail_trainer.train(total_timesteps=100000)
 gail_trainer.gen_algo.save(model_name)