Fixing bugs cropping up because of restructure