Fix context as well