Fix dimension bookkeeping in squeeze