Create TensorDataset

easyFREEDatasetsPyTorch DatasetsPytorch

Problem Statement

Create a TensorDataset from feature and label tensors.

Background

TensorDataset wraps tensors so they can be used with PyTorch's data loading utilities, pairing features with their corresponding labels automatically.

Your Task

Write a function create_tensor_dataset(features, labels) that wraps the two tensors in a TensorDataset and returns it.

Output Format

Return the TensorDataset. The tests check len(dataset), the number of samples it contains.

Example

Input

features=torch.tensor([[1.0, 2.0], [3.0, 4.0]]), labels=torch.tensor([0, 1])

Output

TensorDataset with 2 samples

Reasoning

Inputs: features is a $2 \times 2$ float tensor, one row per sample and one column per feature, so sample 0 is $(1.0, 2.0)$ and sample 1 is $(3.0, 4.0)$ . labels is a length-2 integer tensor holding one class per sample: $0$ and $1$ .
Pairing: both tensors share the same first dimension, $N = 2$ . The dataset lines them up along that dimension, so index $i$ means "row $i$ of features together with entry $i$ of labels".
Worked example: index $0$ gives the pair $\big((1.0, 2.0),\ 0\big)$ and index $1$ gives $\big((3.0, 4.0),\ 1\big)$ . No values are copied, changed or reordered.
Output: the length of the dataset is the number of samples $N$ , not the total number of feature values ( $2 \times 2 = 4$ ) and not the number of features per sample. That is why the output reads TensorDataset with 2 samples.
Notice: a test with 10 samples of 5 features each has length $10$ ; the feature width never affects the count.

Constraints

Use torch.utils.data.TensorDataset
Both tensors will have the same first dimension
Features is a 2D tensor, labels is a 1D tensor

Learn this in

PyTorch FundamentalsFoundations Introduction to MLML 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

TensorDataset is a core PyTorch utility from torch.utils.data that pairs feature tensors (inputs) with label tensors (targets) for supervised learning. It enables seamless integration with DataLoader for batching, shuffling, and parallel loading.

Key prerequisites:

PyTorch tensors: Multi-dimensional arrays with GPU support. Features are typically shape (N, F) where $N$ = samples, $F$ = features. Labels are shape (N,).
Dataset abstraction: PyTorch's Dataset class requires len() (returns sample count) and getitem(idx) (returns (features[idx], labels[idx])).
DataLoader pipeline: TensorDataset → DataLoader(dataset, batch_size=32, shuffle=True) → model training loop.

TensorDataset automatically handles indexing: dataset[i] returns tuple (feature_row_i, label_i).

2. Algorithm Approach

No complex algorithm needed—this is a wrapper construction problem:

Instantiate TensorDataset(features, labels)
Return len(dataset) (equivalent to features.shape due to matching first dimensions)

Mathematical guarantee: Constraints ensure features.shape == labels.shape, so dataset length is well-defined as $N = |$ \text{features} $|_0$ .

3. Step-by-Step Strategy

import torch
from torch.utils.data import TensorDataset

def create_dataset(features, labels):
 # Step 1: Create TensorDataset (pairs tensors automatically)
 dataset = TensorDataset(features, labels)

 # Step 2: Return length (number of samples)
 return len(dataset) # or features.shape

Verification steps:

# Test shape matching (optional, but good practice)
assert features.shape == labels.shape, "Dimension mismatch"

# Usage example
features = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) # (3, 2)
labels = torch.tensor([0, 1, 0]) # (3,)
length = create_dataset(features, labels) # Returns 3

Full training pipeline context:

dataset = TensorDataset(features, labels)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True)
for batch_features, batch_labels in dataloader:
 # batch_features: (2, 2), batch_labels: (2,)
 outputs = model(batch_features)
 loss = criterion(outputs, batch_labels)

4. Common Pitfalls

Import error: Missing from torch.utils.data import TensorDataset
Shape mismatch: Assuming tensors align despite constraints—always verify features.shape == labels.shape
Tensor type: Input must be torch.Tensor, not NumPy arrays (convert with torch.tensor())
Returning dataset instead of length: Function must return int, not TensorDataset object
Device mismatch: Features/labels on different devices (CPU/GPU)—use .to(device)
Empty tensors: features.shape == 0 returns 0 (valid but edge case)

Pitfall	Symptom	Fix
Wrong import	NameError: TensorDataset	from torch.utils.data import TensorDataset
Shape error	RuntimeError: stack expects...	Verify features.shape == labels.shape
Wrong return	Returns object, not int	return len(dataset) not return dataset

5. Time & Space Complexity

Time: $O(1)$

TensorDataset construction: Simple tensor reference (no copying)
len(): Direct shape access

Space: $O(1)$ additional

Stores references to input tensors (no data duplication)
Total memory: $O(N \times F + N)$ where $N$ = samples, $F$ = features

Scaling behavior:

N = 10^6 samples, F = 100 features
- Input memory: ~400 MB (float32)
- TensorDataset overhead: ~8 bytes (tensor pointers)
- DataLoader batches: O(batch_size) temporary memory

This design enables efficient handling of million-sample datasets without memory overhead, making it production-ready for large-scale ML workflows.

solution.py

Python 3.13

Test results