Fixing a small bug with the layout