Update controller
Fix mistake in controller (works better now)