diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py
index cd91bbe..a4afd58 100644
--- a/scratch/etienne/intersimple/gail_options_image.py
+++ b/scratch/etienne/intersimple/gail_options_image.py
@@ -259,8 +259,10 @@ if __name__ == '__main__':
 
     env = NRasterized(**env_settings)
 
-    for _ in sample(env, generator, None, 'low'):
+    for transition in sample(env, generator, None, 'low'):
         env.render()
+        if transition['dones']:
+            break
 
     env.close(filestr='render/'+model_name)