`pred` would be your `p` if I understand correctly.

```python
# Generate a
pred_mask = label_binarize(pred, self.classes_).ravel().astype(bool)
flat_mask = np.logical_and(np.repeat(mask, n_classes), pred_mask)
pred = pred.ravel()
pred[flat_mask] += 1 
```