Fix generator bugs