Small changes to fix a minor bug