Fix for model class generation