What input shape does TransformerEncoder expect in PyTorch?

TransformerEncoder expects input shape (seq_len, batch, d_model) by default. With batch_first=True, it expects (batch, seq_len, d_model). The d_model must match the d_model parameter of the TransformerEncoderLayer. Output shape is the same as input shape. The src_key_padding_mask should have shape (batch, seq_len) as a boolean tensor.

What Input Shape Does TransformerEncoder Expect?

nn.TransformerEncoder expects src of shape (S, N, E), that is (seq_len, batch, d_model), because batch_first defaults to False. Set batch_first=True on the nn.TransformerEncoderLayer and it expects (N, S, E), that is (batch, seq_len, d_model). Unbatched input is (S, E). S is the source sequence length, N the batch size, E is d_model. The output shape always matches the input shape.

“batch_first (bool) – If True, then the input and output tensors are provided as (batch, seq, feature). Default: False (seq, batch, feature).”

PyTorch documentation, torch.nn.TransformerEncoderLayer. Verified 15 August 2026.

What batch_first defaults to in nn.TransformerEncoderLayer

batch_first defaults to False in torch.nn.TransformerEncoderLayer, so PyTorch expects sequence-first tensors of shape (S, N, E) unless you pass batch_first=True to the layer constructor. One detail the signature makes easy to miss. batch_first belongs to nn.TransformerEncoderLayer, not to nn.TransformerEncoder. The stack inherits whatever the layer was built with.

Expected tensor shapes, from the torch.nn.Transformer Shape section
Argument	batch_first=False (default)	batch_first=True	Unbatched
src	(S, N, E)	(N, S, E)	(S, E)
output	(S, N, E)	(N, S, E)	(S, E)
src_mask	(S, S)	(S, S)	(S, S)
src_key_padding_mask	(N, S)	(N, S)	(S)

Shapes above come from the torch.nn.Transformer Shape section, which both nn.TransformerEncoder and nn.TransformerEncoderLayer defer to.

Default Mode (seq_len first)

import torch
import torch.nn as nn

encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)

# Input: (seq_len, batch, d_model)
x = torch.randn(100, 32, 512) # 100 tokens, batch 32, 512 features
output = encoder(x)
print(output.shape) # torch.Size([100, 32, 512])

With batch_first=True (Recommended)

encoder_layer = nn.TransformerEncoderLayer(
  d_model=512, nhead=8, batch_first=True
)
encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)

# Input: (batch, seq_len, d_model)
x = torch.randn(32, 100, 512) # batch 32, 100 tokens, 512 features
output = encoder(x)
print(output.shape) # torch.Size([32, 100, 512])

Using Masks

# Padding mask: (batch, seq_len) — True = ignore this position
padding_mask = torch.zeros(32, 100, dtype=torch.bool)
padding_mask[:, 80:] = True # mask out positions 80-99

# Causal mask: (seq_len, seq_len) — for autoregressive models
causal_mask = nn.Transformer.generate_square_subsequent_mask(100)

output = encoder(x,
  mask=causal_mask,
  src_key_padding_mask=padding_mask
)

Common Mistake

# ERROR: batch_first mismatch
encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
# Default: batch_first=False, expects (seq_len, batch, d_model)
x = torch.randn(32, 100, 512) # This is (batch, seq_len, d_model)
output = encoder(x) # WRONG! Will silently produce bad results

# FIX: set batch_first=True or transpose input
x = x.transpose(0, 1) # (100, 32, 512) ✓

Try the Attention Calculator