Update stage4.py