Update Controllers due to latest 2.0.1 changes