Получите вывод преобразователей GPT слоев Huggingface

Question

Я пытаюсь использовать языковую модель GPT и получить веса, которые она присваивает каждому слову в последнем состоянии генерации текста. Моя модель - GPT2 из библиотеки трансформеров. Вот как я называю предварительно обученную модель:

      tokenizer = AutoTokenizer.from_pretrained(
"HooshvareLab/gpt2-fa-poetry"
) 

model = AutoModelForCausalLM.from_pretrained(
    "HooshvareLab/gpt2-fa-poetry"
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model.to(device)

Моя цель - использовать эту информацию из последнего слоя этой модели (матрица с длиной словаря после активации softmax) и использовать ее в сочетании с другой моделью.

Я пытаюсь сделать это в TensorFlowPlease, но поделитесь своими комментариями, если вы считаете, что в PyTorch есть более простые и удобные способы сделать это.