Fixing bugs from Evgeny's feedback