Update masks class to be on correct element