Create Lambda Transform

mediumFREETransformsPyTorch TransformsPytorch

Problem Statement

Create a Lambda transform that converts integer labels to one-hot encoded tensors.

Background

Lambda transforms let you wrap any custom function as a torchvision transform. One-hot encoding is a common technique that converts a class label into a binary vector — useful for classification tasks.

Your Task

Write a function create_onehot_transform(num_classes) that returns a Lambda transform. When applied to an integer label, the transform should produce a dictionary with the one-hot tensor and its sum.

Output Format

The transform should return a dictionary with keys: "tensor" (one-hot list), "sum" (float, should be 1.0).

Example

Input

num_classes=5, then apply to label=2

Output

{"tensor": [0.0, 0.0, 1.0, 0.0, 0.0], "sum": 1.0}

Reasoning

Inputs: num_classes $= 5$ sets the vector length: one slot per class, indexed $0$ to $4$ . The label 2 is this example's class.
Math: the one-hot vector $e$ for label $k$ has $e_i = 1$ when $i = k$ and $e_i = 0$ otherwise. With $k = 2$ : $e_0 = 0,\ e_1 = 0,\ e_2 = 1,\ e_3 = 0,\ e_4 = 0$ , which gives [0.0, 0.0, 1.0, 0.0, 0.0].
Sum: $\sum_{i=0}^{4} e_i = 0 + 0 + 1 + 0 + 0 = 1$ , so sum is 1.0. It is $1$ for any valid label, since exactly one slot is lit.
Output: a dictionary: tensor is the vector as a plain list of floats, sum is that total. Both are floats ( $1.0$ , not $1$ ): the vector is a probability distribution with all its mass on one class.
Notice: only the position of the $1$ depends on the label. Label 0 with $3$ classes gives [1.0, 0.0, 0.0], and label $3$ with $4$ classes gives [0.0, 0.0, 0.0, 1.0], the last slot.

Constraints

Use torchvision.transforms.Lambda
Use torch.zeros and scatter_ for one-hot encoding
num_classes will be between 2 and 10
Labels will be valid integers from 0 to num_classes-1

Learn this in

NumPy & Data WranglingFoundations ClassificationML Study Plan Text RepresentationNLP Study Plan

More from Pytorch

Apply Softmax for Predictionseasy Count Model Parametersmedium Build with nn.Sequentialeasy Use nn.Flatten Layereasy Create a Simple Neural Network Classeasy

Concept primer

1. Background Knowledge

One-Hot Encoding converts categorical labels into binary vectors where only the target class index is 1, others are 0. For $C$ classes and label $k$ , the encoding is: $\mathbf{e}_k = [0, 0, \dots, 1, \dots, 0] \quad \text{where } e_k[i] = \delta_{ik}$

torchvision.transforms.Lambda wraps any Python callable as a transform for datasets/pipelines:

transform = Lambda(lambda x: torch.zeros(5).scatter_(0, x, 1.0))

torch.zeros(dim).scatter_(dim, index, value) efficiently creates one-hot encodings:

Creates zero tensor of size dim
Sets tensor[index] = value (in-place)

Key Prerequisites:

PyTorch tensor basics (torch.zeros, indexing)
Understanding of classification data pipelines
Familiarity with torchvision.transforms

2. Algorithm Approach

Core Algorithm: Scatter-based One-Hot

1. Create zero tensor: torch.zeros(num_classes, dtype=torch.float)
2. Set target position: tensor.scatter_(0, label, 1.0)
3. Compute checksum: tensor.sum()
4. Package as dict: {"tensor": tensor.tolist(), "sum": sum_value}

Why scatter_?

O(C) time (constant per class count)
Vectorized, GPU-compatible
In-place (memory efficient)
Handles batching automatically

Mathematical Foundation: $\text{onehot}(x, C) = \sum_{i=0}^{C-1} \delta_{xi} \cdot \mathbf{e}_i$

3. Step-by-Step Strategy

def create_onehot_transform(num_classes):
 def onehot_fn(label):
 tensor = torch.zeros(num_classes) # Step 1: Zero tensor
 tensor.scatter_(0, torch.tensor(label), 1.0) # Step 2: Set 1 at index
 return {
 "tensor": tensor.tolist(), # Step 3: Convert to list
 "sum": tensor.sum().item() # Step 4: Compute sum as float
 }
 return Lambda(onehot_fn) # Step 5: Wrap in Lambda

Execution Flow:

label=2, num_classes=5
↓
torch.zeros(5) → [0., 0., 0., 0., 0.]
↓
scatter_(0, tensor(2), 1.0) → [0., 0., 1., 0., 0.]
↓
{"tensor": [0.0, 0.0, 1.0, 0.0, 0.0], "sum": 1.0}

4. Common Pitfalls

Pitfall	Problem	Fix
tensor.sum() returns tensor	Dict expects float, not tensor	Use .item(): tensor.sum().item()
Return tensor directly	Must return dict with exact keys	{"tensor": tensor.tolist(), "sum": 1.0}
Use torch.eye() or indexing	Inefficient, not using scatter_	tensor.scatter_(0, label, 1.0)
Forget torch.tensor(label)	scatter_ needs tensor index	torch.tensor(label, dtype=torch.long)
Integer tensor dtype	Output must be float [0.0, 1.0]	torch.zeros(..., dtype=torch.float)
tensor.cpu() or .detach()	Unnecessary overhead	Raw tensor is fine

Testing Edge Cases:

transform(0) # First class: [1.0, 0.0,...]
transform(num_classes-1) # Last class

5. Time & Space Complexity

Time Complexity: $O(C)$ per transform call

torch.zeros(C): $O(C)$
scatter_(0, idx, 1.0): $O(1)$ (scatter is constant time)
tolist(): $O(C)$
Total: $O(C)$ where $C =$ num_classes (2-10)

Space Complexity: $O(C)$

Single tensor of size $C$
Output dict references same tensor (no copy)

Scalability Notes:

Batch of N labels: O(N × C) → Embarrassingly parallel
GPU: Native PyTorch ops → Automatic acceleration
C ≤ 10 constraint → Negligible overhead in pipelines

Complete Solution:

import torch
from torchvision.transforms import Lambda

def create_onehot_transform(num_classes):
 def onehot_fn(label):
 tensor = torch.zeros(num_classes, dtype=torch.float)
 tensor.scatter_(0, torch.tensor(label), 1.0)
 return {
 "tensor": tensor.tolist(),
 "sum": tensor.sum().item()
 }
 return Lambda(onehot_fn)

This creates a production-ready, efficient transform ready for any PyTorch Dataset/DataLoader pipeline.

solution.py

Python 3.13

Test results