Updated to fix some issues