Create TensorDataset

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:

1. Instantiate TensorDataset(features, labels)
2. 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)

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.