Rescheduled output stage to use smaller blocks on GPU.