import torch

def vit_cls_token_mean(x: torch.Tensor):
    # Mean class token
    return x[:, 0, :].mean(dim=0)

def vit_img_tokens_mean(x: torch.Tensor):
    # Mean image token
    return x[:, 1:, :].mean(dim=(0, 1))
