PyTorch · Class

nn.Embedding

nn.Embedding is a trainable lookup table for a fixed dictionary of vectors. Integer indices retrieve the corresponding rows, commonly converting token IDs into dense representations.

How it works

Embedding owns a learnable matrix with one row per dictionary entry and one column per coordinate. Calling it with integer indices selects the corresponding rows, preserving the input shape and appending the embedding dimension; optional settings can freeze padding rows, constrain row norms, rescale frequent-token gradients, or request sparse weight gradients.

Syntax

torch.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None, _freeze=False, device=None, dtype=None)

Parameters

Name	Type	Description
`num_embeddings`	`int`	Number of embedding vectors in the dictionary.
`embedding_dim`	`int`	Number of values in each embedding vector.
`padding_idx`	`int, optional`	Optional row excluded from gradients and used as padding.
`max_norm`	`float, optional`	Optional maximum row norm.
`norm_type`	`float, optional`	p-norm used with max_norm; default 2.
`scale_grad_by_freq`	`bool, optional`	Scale gradients by inverse minibatch frequency; default False.
`sparse`	`bool, optional`	If True, produce a sparse gradient for weight.
`_weight`	`Tensor, optional`	Optional initial weight tensor.
`_freeze`	`bool, optional`	Optional flag controlling freezing; default False.
`device`	`torch.device, optional`	Optional device for the module.
`dtype`	`torch.dtype, optional`	Optional data type for the module.

Returns

A tensor containing the embedding vectors selected by the input indices; the output preserves the input's leading shape and appends embedding_dim.

Key points

The class represents a fixed-size dictionary of learnable embedding vectors.
Forward evaluation maps each supplied index to its corresponding vector.
num_embeddings specifies how many rows the table contains.
embedding_dim specifies the width of every table row.
max_norm caps oversized row norms during lookup.
scale_grad_by_freq reduces the gradient contribution of tokens that occur more often in a minibatch.

Common mistakes

Inputs must use an integer tensor type, conventionally torch.long; floating-point indices are invalid.
With sparse=True, the weight gradient is sparse, so the optimizer must support sparse gradients.
padding_idx starts as a zero vector in a newly created module, although it can later be assigned another padding value.

Example: Lookup embedding vectors

import torch
from torch import nn

embedding = nn.Embedding(num_embeddings=4, embedding_dim=3, padding_idx=0)
indices = torch.tensor([0, 2, 3], dtype=torch.long)
output = embedding(indices)
print(output.shape)  # torch.Size([3, 3])

Creates a four-entry table of three-dimensional vectors and retrieves the rows selected by integer indices; index 0 is reserved as padding.

When to use it

Token and word representations

Convert token IDs into dense vectors before an NLP model.

Categorical feature embeddings

Represent categorical IDs such as users, products, or locations.

Padded sequence inputs

Handle variable-length padded sequences while keeping padding fixed.

Interview questions

QWhat does nn.Embedding do?

It is a trainable lookup table with shape (num_embeddings, embedding_dim). The input supplies integer row indices, and the output contains the selected embedding vectors.

Follow-up: How does it differ from EmbeddingBag? Embedding returns the selected vectors directly, whereas EmbeddingBag pools groups of embeddings without materializing the intermediate vectors.

QHow does padding_idx affect training?

padding_idx marks a row that does not contribute to the gradient, so training does not update that row. A newly constructed module initializes this row to zeros, though it may be changed to another padding vector.

Follow-up: Why is this useful? It lets padded sequence positions share a stable representation while other token rows remain trainable.

QWhat do max_norm and scale_grad_by_freq control?

max_norm renormalizes rows whose norm exceeds the limit, using norm_type to choose the p-norm. scale_grad_by_freq instead scales gradients by the inverse frequency of each word in the minibatch.

Follow-up: What does sparse=True change? It makes gradients with respect to the weight matrix sparse, restricting compatible optimizers.

torch.nn.EmbeddingBag
torch.nn.modules.sparse

Source: Embedding — PyTorch 2.13 documentation. Checked 2026-10-16.