Code cleaned up of previous patch