50. Tiny Text Classifier with a Transformer Intro in PyTorch

50.3 Mean-Embed Classifier

By Ravindra Bagale · Last updated: 24 October 2026

Why. A minimal net shows embeddings working before full multi-head attention code.

What. Embedding → mask pads → mean pool → Linear to 1 logit. (Stretch: replace mean pool with nn.TransformerEncoder later.)

How.

import torch
import torch.nn as nn

class MeanEmbedClassifier(nn.Module):
 def __init__(self, vocab_size=8, dim=16):
 super().__init__()
 self.emb = nn.Embedding(vocab_size, dim, padding_idx=0)
 self.fc = nn.Linear(dim, 1)
 def forward(self, token_ids):
 x = self.emb(token_ids) # (B, T, D)
 mask = (token_ids != 0).float().unsqueeze(-1)
 summed = (x * mask).sum(dim=1)
 counts = mask.sum(dim=1).clamp(min=1.0)
 return self.fc(summed / counts)

print(MeanEmbedClassifier())