Slight modification to controller