---
name: pytorch-patterns
description: Use when writing PyTorch code to follow best practices and common patterns
---

# PyTorch Patterns Skill

## Overview

This skill provides best practices and common patterns for writing PyTorch code. Use this when implementing neural networks, training loops, data loading, and related deep learning infrastructure.

**Announce:** "I'm using the pytorch-patterns skill for best practice code."

---

## Model Definition

### Basic nn.Module Pattern

```python
from typing import NamedTuple
import torch
import torch.nn as nn
import torch.nn.functional as F


class ModelConfig:
    """Configuration for the model."""
    def __init__(
        self,
        input_dim: int = 768,
        hidden_dim: int = 512,
        output_dim: int = 10,
        dropout: float = 0.1,
        num_layers: int = 2,
    ):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.dropout = dropout
        self.num_layers = num_layers


class ModelOutput(NamedTuple):
    """Typed output container for model forward pass."""
    logits: torch.Tensor
    hidden_states: torch.Tensor
    attention_weights: torch.Tensor | None = None


class MyModel(nn.Module):
    """Example model following best practices."""
    
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        
        # Define layers
        self.input_proj = nn.Linear(config.input_dim, config.hidden_dim)
        self.layers = nn.ModuleList([
            nn.Linear(config.hidden_dim, config.hidden_dim)
            for _ in range(config.num_layers)
        ])
        self.output_proj = nn.Linear(config.hidden_dim, config.output_dim)
        self.dropout = nn.Dropout(config.dropout)
        self.layer_norm = nn.LayerNorm(config.hidden_dim)
        
        # Initialize weights
        self._init_weights()
    
    def _init_weights(self):
        """Initialize model weights with appropriate schemes."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)
    
    @property
    def device(self) -> torch.device:
        """Get the device of model parameters."""
        return next(self.parameters()).device
    
    def forward(self, x: torch.Tensor) -> ModelOutput:
        """Forward pass with typed output."""
        # Input projection
        hidden = self.input_proj(x)
        hidden = self.layer_norm(hidden)
        hidden = F.gelu(hidden)
        
        # Process through layers
        for layer in self.layers:
            residual = hidden
            hidden = layer(hidden)
            hidden = self.dropout(hidden)
            hidden = F.gelu(hidden)
            hidden = hidden + residual  # Residual connection
        
        # Output projection
        logits = self.output_proj(hidden)
        
        return ModelOutput(
            logits=logits,
            hidden_states=hidden,
            attention_weights=None,
        )
```

---

## Device Management

### Device Property Pattern

```python
class DeviceAwareModule(nn.Module):
    """Module with device awareness."""
    
    @property
    def device(self) -> torch.device:
        """Infer device from model parameters."""
        return next(self.parameters()).device
    
    @property
    def dtype(self) -> torch.dtype:
        """Infer dtype from model parameters."""
        return next(self.parameters()).dtype
```

### Training Script Device Setup

```python
def get_device() -> torch.device:
    """Get the best available device."""
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Using CUDA: {torch.cuda.get_device_name(0)}")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("Using Apple MPS")
    else:
        device = torch.device("cpu")
        print("Using CPU")
    return device


# Usage
device = get_device()
model = MyModel(config).to(device)
```

---

## Training Loop

### Standard Training Epoch

```python
def train_epoch(
    model: nn.Module,
    dataloader: torch.utils.data.DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    max_grad_norm: float | None = 1.0,
) -> dict[str, float]:
    """Train for one epoch."""
    model.train()
    total_loss = 0.0
    num_batches = 0
    
    for batch_idx, (inputs, targets) in enumerate(dataloader):
        # Move to device
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        # Forward pass
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs.logits, targets)
        
        # Backward pass
        loss.backward()
        
        # Gradient clipping (optional but recommended)
        if max_grad_norm is not None:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
        
        # Update weights
        optimizer.step()
        
        # Accumulate metrics (use .item() to prevent memory leak)
        total_loss += loss.item()
        num_batches += 1
    
    return {
        "train_loss": total_loss / num_batches,
    }
```

### Evaluation Function

```python
@torch.no_grad()
def evaluate(
    model: nn.Module,
    dataloader: torch.utils.data.DataLoader,
    criterion: nn.Module,
    device: torch.device,
) -> dict[str, float]:
    """Evaluate the model."""
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, targets in dataloader:
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        outputs = model(inputs)
        loss = criterion(outputs.logits, targets)
        
        total_loss += loss.item()
        
        # Calculate accuracy
        predictions = outputs.logits.argmax(dim=-1)
        correct += (predictions == targets).sum().item()
        total += targets.size(0)
    
    return {
        "eval_loss": total_loss / len(dataloader),
        "accuracy": correct / total,
    }
```

---

## Data Loading

### Custom Dataset Pattern

```python
from torch.utils.data import Dataset, DataLoader


class CustomDataset(Dataset):
    """Example custom dataset."""
    
    def __init__(self, data: list, transform=None):
        self.data = data
        self.transform = transform
    
    def __len__(self) -> int:
        return len(self.data)
    
    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        item = self.data[idx]
        features = torch.tensor(item["features"], dtype=torch.float32)
        label = torch.tensor(item["label"], dtype=torch.long)
        
        if self.transform:
            features = self.transform(features)
        
        return features, label
```

### DataLoader Helper

```python
def get_dataloader(
    dataset: Dataset,
    batch_size: int = 32,
    shuffle: bool = True,
    num_workers: int = 4,
    pin_memory: bool = True,
    drop_last: bool = False,
) -> DataLoader:
    """Create a DataLoader with sensible defaults."""
    # Adjust for Windows/macOS compatibility
    if num_workers > 0:
        import platform
        if platform.system() == "Windows":
            num_workers = 0  # Windows has issues with multiprocessing
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory and torch.cuda.is_available(),
        drop_last=drop_last,
        persistent_workers=num_workers > 0,
    )
```

---

## Checkpointing

### Save Checkpoint

```python
def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    epoch: int,
    loss: float,
    path: str,
    scheduler: torch.optim.lr_scheduler._LRScheduler | None = None,
    **kwargs,
) -> None:
    """Save a training checkpoint."""
    checkpoint = {
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }
    
    if scheduler is not None:
        checkpoint["scheduler_state_dict"] = scheduler.state_dict()
    
    # Add any additional data
    checkpoint.update(kwargs)
    
    torch.save(checkpoint, path)
    print(f"Checkpoint saved to {path}")
```

### Load Checkpoint

```python
def load_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
    scheduler: torch.optim.lr_scheduler._LRScheduler | None = None,
    device: torch.device | None = None,
) -> dict:
    """Load a training checkpoint."""
    # Use weights_only=True for security (prevents arbitrary code execution)
    checkpoint = torch.load(path, map_location=device, weights_only=True)
    
    model.load_state_dict(checkpoint["model_state_dict"])
    
    if optimizer is not None and "optimizer_state_dict" in checkpoint:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    
    if scheduler is not None and "scheduler_state_dict" in checkpoint:
        scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
    
    print(f"Loaded checkpoint from {path} (epoch {checkpoint.get('epoch', 'unknown')})")
    return checkpoint
```

---

## Reproducibility

### Set Seed Function

```python
import random
import numpy as np
import torch


def set_seed(seed: int = 42, deterministic: bool = True) -> None:
    """Set random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)  # For multi-GPU
    
    if deterministic:
        # Makes operations deterministic but may reduce performance
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    else:
        # Better performance but non-deterministic
        torch.backends.cudnn.deterministic = False
        torch.backends.cudnn.benchmark = True


# Usage at start of training script
set_seed(42)
```

---

## Common Gotchas

### 1. In-Place Operations Breaking Autograd

```python
# BAD - In-place operation can break autograd graph
def bad_forward(x):
    x += 1  # In-place modification
    return x * 2

# GOOD - Create new tensor
def good_forward(x):
    x = x + 1  # Creates new tensor
    return x * 2
```

### 2. Memory Leaks from Not Detaching

```python
# BAD - Keeps computation graph in memory
losses = []
for batch in dataloader:
    loss = model(batch)
    losses.append(loss)  # Holds entire graph!

# GOOD - Detach with .item() for scalars
losses = []
for batch in dataloader:
    loss = model(batch)
    losses.append(loss.item())  # Just the number, no graph

# GOOD - Detach for non-scalar tensors
features = []
for batch in dataloader:
    feat = model.encode(batch)
    features.append(feat.detach().cpu())  # Detached, moved to CPU
```

### 3. Mixed Precision Training

```python
from torch.cuda.amp import autocast, GradScaler


def train_with_amp(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
) -> float:
    """Training with Automatic Mixed Precision."""
    model.train()
    scaler = GradScaler()
    total_loss = 0.0
    
    for inputs, targets in dataloader:
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        optimizer.zero_grad()
        
        # Forward pass with autocast
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs.logits, targets)
        
        # Backward pass with scaled gradients
        scaler.scale(loss).backward()
        
        # Unscale before clipping
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        
        # Step with scaler
        scaler.step(optimizer)
        scaler.update()
        
        total_loss += loss.item()
    
    return total_loss / len(dataloader)
```

### 4. Forgetting model.eval() and model.train()

```python
# BAD - Dropout/BatchNorm behave incorrectly
def evaluate_bad(model, dataloader):
    # model.train() is still active!
    for batch in dataloader:
        output = model(batch)

# GOOD - Always set mode explicitly
def evaluate_good(model, dataloader):
    model.eval()  # Disable dropout, use running stats for BatchNorm
    with torch.no_grad():
        for batch in dataloader:
            output = model(batch)
    model.train()  # Restore training mode if continuing
```

### 5. Proper Tensor Creation on Device

```python
# BAD - Creates on CPU then moves (slow)
tensor = torch.zeros(100, 100).to(device)

# GOOD - Create directly on device
tensor = torch.zeros(100, 100, device=device)

# GOOD - Create with same device/dtype as reference
tensor = torch.zeros_like(reference_tensor)
tensor = torch.empty(100, 100, device=model.device, dtype=model.dtype)
```

---

## Quick Reference Checklist

- [ ] Use `nn.Module` with proper `__init__` and `forward`
- [ ] Initialize weights with `_init_weights` method
- [ ] Use `@property device` for device inference
- [ ] Always use `.item()` when logging scalar losses
- [ ] Use `@torch.no_grad()` decorator for evaluation
- [ ] Call `model.train()` and `model.eval()` explicitly
- [ ] Use `weights_only=True` when loading checkpoints
- [ ] Set seeds for reproducibility at script start
- [ ] Avoid in-place operations in forward pass
- [ ] Use mixed precision for faster training on GPU
- [ ] Clip gradients to prevent exploding gradients
- [ ] Use `pin_memory=True` with CUDA for faster data transfer