Trying to fix a few bugs