Fixed small bug in viewController