Fix th logic from the previous commit.