Update all_stages.py