src/transformers/models/flava/modeling_flava.py/FlavaImageCodebook/get_codebook_probs
class FlavaImageCodebook:    def get_codebook_probs(self, pixel_values: torch.Tensor) -> torch.Tensor:
        z_logits = self.blocks(pixel_values)
        return nn.Softmax(dim=1)(z_logits)