Minor modification on controller