def METHOD_NAME(self, head: str, tokens: torch.LongTensor, return_logits: bool = False):
    if tokens.dim() == 1:
        tokens = tokens.unsqueeze(0)
    features = self.extract_features(tokens.to(device=self.device))
    sentence_representation = features[
        tokens.eq(self.task.source_dictionary.eos()), :
    ].view(features.size(0), -1, features.size(-1))[:, -1, :]
    logits = self.model.classification_heads[head](sentence_representation)
    if return_logits:
        return logits
    return F.log_softmax(logits, dim=-1)