Update base_model.py