Clean up
Fixing bug with generator