Update drop_dims, now works with more than 2 dims passed