Fixing up a couple bugs