PyTorch ships the whole transformer block. nn.TransformerEncoderLayer is multi-head attention and a small MLP, each with a residual connection and layer norm; nn.TransformerEncoder stacks copies of it.

Attention itself has no idea of order. Shuffle the positions of the input and the output is shuffled the same way, nothing else changes:

import torch
from torch import nn

torch.manual_seed(0)
layer = nn.TransformerEncoderLayer(
    d_model=8, nhead=2, dim_feedforward=16,
    dropout=0.0, batch_first=True)
x = torch.randn(1, 4, 8)
perm = torch.tensor([2, 0, 3, 1])
a = layer(x)[:, perm]
b = layer(x[:, perm])
print(torch.allclose(a, b, atol=1e-5))

So "the cat chased the dog" and "the dog chased the cat" look alike. Order enters only through position embeddings: a learned vector per position, added to each token's embedding before the first layer.

Set dropout=0.0 for small experiments; the default of 0.1 adds noise that only pays off on large data.

Using the given order_data and train, write OrderClassifier(vocab, d, context): self.tokens = nn.Embedding(vocab, d), self.positions = nn.Embedding(context, d), self.encoder, an nn.TransformerEncoder of one nn.TransformerEncoderLayer(d, 2, 2 * d, dropout=0.0, batch_first=True), and self.head = nn.Linear(d, 2). forward(ids) embeds tokens plus positions, encodes, averages over positions and returns the head's two scores per sequence.

The explanation opens here once every check passes.