PyTorch Tensors — Silent CPU Fallback Kills GPU Utilization

PyTorch Tensors are the fundamental data structure in PyTorch — every input, weight, gradient, and output is a tensor. They are multidimensional arrays that mirror NumPy's API but add two capabilities that NumPy does not have: GPU acceleration via CUDA and automatic differentiation via Autograd.

The key design decision: tensors are not just data containers. When requires_grad=True, they become nodes in a dynamic computation graph. Every operation on them is recorded as the forward pass executes, enabling automatic gradient computation when you call .backward(). This is what makes neural network training tractable — without it, you would manually compute partial derivatives for every parameter on every update, which is not realistic at any modern model size.

The production failure pattern: device mismatch. A tensor on CPU cannot participate in the same operation as a tensor on GPU. PyTorch raises RuntimeError: Expected all tensors to be on the same device immediately and clearly. What is not clear is which tensor is on the wrong device. The fix is always the same: establish device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') once, pass it to every tensor creation call, and add an assertion at training start that verifies model parameters and input data are on the same device.

As of 2026, there is a third dimension worth knowing: Apple Silicon. PyTorch supports MPS (Metal Performance Shaders) on M-series Macs via device='mps', which provides meaningful GPU acceleration on MacBooks without CUDA. The same .to(device) pattern applies — the device string changes, the code does not.

What Is a PyTorch Tensor and Why Does It Exist?

A PyTorch tensor is a multidimensional array that was designed to solve a problem NumPy cannot: running massive parallel computations on a GPU while simultaneously tracking every operation for automatic gradient computation.

NumPy arrays are excellent for scientific computing — fast, well-documented, universally supported. But they have two hard limits. First, they run only on CPU. Second, they have no concept of a computation graph. This means that if you want to train a neural network with NumPy, you implement backpropagation manually — computing partial derivatives by hand for every layer, every parameter, every batch. That is tractable for a two-layer toy network and completely unworkable for anything beyond it.

PyTorch tensors solve both problems. The storage layer underneath a tensor can live on CPU, on an NVIDIA GPU via CUDA, or on Apple Silicon via MPS. When the storage is on GPU, every matrix operation dispatches to CUDA kernels that execute in parallel across thousands of GPU cores — this is why large matrix multiplications are 10–100x faster on GPU for the sizes neural networks operate on. When requires_grad=True, the tensor records every operation as part of a dynamic computation graph. When .backward() is called on the loss, that graph is traversed in reverse and .grad is filled in on every participating tensor via the chain rule — one backward pass, all gradients computed simultaneously.

The architectural detail worth understanding: a tensor is a view into a storage object. The tensor knows its shape, stride, dtype, and device. The storage holds the raw bytes. Operations like .transpose() and .permute() create new tensor views with reordered strides without moving any data in memory — the storage stays identical. This is efficient but has a consequence: the resulting tensor is non-contiguous, and .view() will refuse to work on it because .view() requires elements to be laid out in memory in the same order they are addressed logically. The fix is .contiguous(), which copies the data into a new storage with the expected memory order.

As of PyTorch 2.x, there is a fourth dimension: compiled tensors. torch.compile traces your forward pass and compiles it into optimised kernels using TorchInductor. The tensor API is identical — you add one decorator and the same tensor operations run in a fused, optimised form. For production inference workloads in 2026, torch.compile is the highest-leverage single change you can make to a trained model.

import torch

# Device selection — this pattern belongs at the top of every script
# Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU fallback
if torch.cuda.is_available():
    device = torch.device("cuda")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
    device = torch.device("mps")  # Apple Silicon GPU — supported from PyTorch 1.12+
else:
    device = torch.device("cpu")

print(f"Using device: {device}")


def initialize_forge_tensors():
    # Creating tensors from Python data — torch.tensor (lowercase) infers dtype
    data = [[1, 2], [3, 4]]
    x_data = torch.tensor(data, dtype=torch.float32, device=device)
    print(f"From list — shape: {x_data.shape}, device: {x_data.device}, dtype: {x_data.dtype}")

    # Creating tensors with specific values — pass device at creation, not after
    x_ones  = torch.ones(2, 3, dtype=torch.float32, device=device)
    x_zeros = torch.zeros(2, 3, dtype=torch.float32, device=device)
    x_rand  = torch.randn(2, 3, device=device)  # standard normal distribution
    print(f"\nones:\n{x_ones}")
    print(f"zeros:\n{x_zeros}")
    print(f"randn:\n{x_rand}")

    # requires_grad=True — opts this tensor into the computation graph
    # Only set this on learnable parameters, never on input data
    x_grad = torch.randn(3, 3, requires_grad=True, device=device)
    print(f"\nGradient tracking enabled: {x_grad.requires_grad}")
    print(f"grad_fn before operation: {x_grad.grad_fn}")  # None — leaf tensor

    # Any operation on a requires_grad tensor produces a non-leaf tensor with grad_fn
    y = x_grad ** 2  # element-wise square
    print(f"grad_fn after operation:  {y.grad_fn}")  # PowBackward0

    return x_data, x_grad


x_data, x_grad = initialize_forge_tensors()

# Tensor metadata inspection — useful diagnostic at the start of debugging
for name, t in [("x_data", x_data), ("x_grad", x_grad)]:
    print(f"{name}: shape={t.shape}, dtype={t.dtype}, device={t.device}, requires_grad={t.requires_grad}")

Enterprise Data Pipelines: SQL to Tensor Conversion

In production ML systems, training data rarely arrives as a Python list. It lives in a relational database — normalised, versioned, and filtered by business logic before it ever reaches a tensor. The conversion path from SQL to tensor has two meaningful implementation choices that trade memory efficiency against safety, and getting this wrong at scale causes OOM crashes or silent data corruption.

The recommended pipeline: query SQL into a Pandas DataFrame or directly into a NumPy array via the database cursor's fetchall method. Then convert to a tensor using torch.from_numpy(arr) for zero-copy conversion — the tensor and the NumPy array share the same underlying memory, so no data is duplicated. Move to GPU with .to(device). This entire path keeps memory usage as low as possible for datasets that fit in RAM.

The danger with torch.from_numpy(): because the tensor and the NumPy array share memory, modifying the NumPy array after conversion will silently change the tensor's data. In a pipeline where the DataFrame is reused or mutated for other purposes, this can corrupt training data without any error. If the NumPy array may change after conversion, use torch.tensor(arr, device=device) instead — it creates an independent copy at the cost of a second allocation.

For datasets that do not fit in RAM — anything beyond a few hundred thousand rows with high-dimensional features — loading everything at once causes OOM before training starts. The correct pattern is a custom Dataset that queries or reads one batch at a time in __getitem__, combined with a DataLoader that parallelises the fetching. This keeps memory usage proportional to batch size regardless of dataset size.

-- Extracting normalised features for tensor conversion
-- We pull only the columns needed for training — not SELECT *
-- The WHERE clause filters to verified samples only, matching the Dataset class expectation
-- LIMIT controls batch size for incremental loading patterns

SELECT
    feature_a,
    feature_b,
    target_label
FROM io.thecodeforge.analytics_table
WHERE status = 'processed'
  AND is_verified = TRUE
  AND split_tag = 'train'
ORDER BY sample_id ASC
LIMIT 10000;

-- For incremental loading in a custom Dataset.__getitem__:
-- Use OFFSET :offset LIMIT :batch_size with parameterised queries
-- Never load the full table in __init__ — keep memory proportional to batch size

Standardising Environments with Docker

PyTorch's GPU support depends on a precise version compatibility chain: host NVIDIA driver → CUDA runtime → cuDNN → PyTorch. A mismatch at any point produces a silent failure — torch.cuda.is_available() returns False, tensors silently stay on CPU, and training runs at a fraction of expected speed with no error message to guide you. Docker solves this by fixing the entire stack in one image tag.

The compatibility rule: the CUDA version in the Docker base image must be less than or equal to the CUDA version supported by the host machine's NVIDIA driver. The driver's maximum supported CUDA version is shown in the top-right corner of nvidia-smi output. If the image requests a higher CUDA version than the driver supports, PyTorch loads but CUDA initialisation fails silently. The fix is always to pick a base image whose CUDA version is at or below what nvidia-smi reports.

Two environment variables matter for GPU-enabled containers. NVIDIA_VISIBLE_DEVICES controls which physical GPUs the container can see — set it to all for training containers or to a specific index when you need to isolate workloads on a multi-GPU host. NVIDIA_DRIVER_CAPABILITIES tells the NVIDIA Container Toolkit which driver features to expose — compute gives you CUDA compute, utility gives you nvidia-smi inside the container. Both should be set in the Dockerfile rather than passed at runtime so they are reproducible.

The verification step that should run before every training job: add a startup script that calls torch.cuda.is_available(), prints the GPU name and total memory, and exits with a non-zero code if CUDA is not available when it is expected. This turns silent CPU fallback into an immediate and obvious failure that stops the job before it wastes hours of compute.

# Pin specific PyTorch and CUDA versions — never use 'latest' in production
# 'latest' changes silently and makes training runs non-reproducible
# Check pytorch.org for the full compatibility matrix before changing these
FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime

WORKDIR /app

# Expose all GPUs to the container runtime
# NVIDIA_VISIBLE_DEVICES=all makes every physical GPU available
# NVIDIA_DRIVER_CAPABILITIES=compute,utility enables CUDA compute and nvidia-smi
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

# Prevent CUDA memory fragmentation on long training runs
# Without this, OOM errors can occur even when total free memory is sufficient
ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

# Startup verification — catches CUDA misconfiguration before wasting compute
# Remove the CUDA assertion for CPU-only deployment targets
RUN python -c "
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'GPU: {torch.cuda.get_device_name(0)}')
    print(f'CUDA version: {torch.version.cuda}')
"

# Run with: docker run --gpus all --shm-size=2g -v /data:/data thecodeforge/torch-runtime:2.3.1
CMD ["python", "ForgeTensorBasics.py"]

Common Mistakes and How to Avoid Them

Most tensor bugs in production trace back to a handful of patterns. Knowing them in advance is the difference between a 30-second fix and a four-hour debugging session.

Device mismatch is the most common runtime crash. The error message is unambiguous — RuntimeError: Expected all tensors to be on the same device — but identifying which tensor is on the wrong device requires checking .device on each one. The usual culprit is a target tensor or loss buffer created inside the training loop without .to(device), while the model output is correctly on GPU. The fix: establish device once and pass it to every tensor creation call in the loop, not just to the model.

torch.Tensor (capital T) versus torch.tensor (lowercase t) is a confusion that trips up developers coming from other frameworks. torch.Tensor is the tensor class constructor — calling torch.Tensor([1, 2, 3]) creates a float32 tensor from data, but it is the long form of torch.FloatTensor and does not perform dtype inference. torch.tensor (lowercase) is the factory function that infers dtype from the input, always creates a copy, accepts device and requires_grad arguments, and is the correct way to create a tensor from data. In production code, always use torch.tensor().

In-place operations on gradient-tracked tensors corrupt the computation graph. When you call a.add_(b), the original value of a — which autograd needs to compute a's gradient during backpropagation — is destroyed. PyTorch raises RuntimeError: a leaf Variable that requires grad is being used in an in-place operation if it catches this immediately, but in some cases the graph is silently corrupted and gradients are wrong without any error. The rule: avoid trailing underscores (add_, mul_, fill_, zero_) on any tensor with requires_grad=True.

Views versus copies is the final major source of confusion. .view(), slicing, and .transpose() all return views that share storage with the original tensor. Modifying a view modifies the original. .clone() creates an independent copy. If you need to modify a slice without affecting the source tensor — common when building augmented versions of a batch — always call .clone() first.

import torch

# Device selection — establish once, pass everywhere
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")

# --- Mistake 1: Device mismatch ---
# WRONG: x is on GPU, y is on CPU — crashes on the addition
# x = torch.ones(5, device='cuda')
# y = torch.zeros(5)  # CPU by default
# z = x + y  # RuntimeError: Expected all tensors to be on the same device

# CORRECT: both created on the same device
x = torch.ones(5, device=device)
y = torch.zeros(5, device=device)
z = x + y
print(f"Device mismatch fixed: z = {z}")

# --- Mistake 2: torch.Tensor vs torch.tensor ---
# WRONG: torch.Tensor (class) — confusing, no dtype inference, no device argument
# confused = torch.Tensor([1, 2, 3])  # always float32, always CPU

# CORRECT: torch.tensor (factory function) — explicit dtype, device, requires_grad
clear = torch.tensor([1, 2, 3], dtype=torch.float32, device=device)
print(f"torch.tensor result: {clear}, device: {clear.device}")

# --- Mistake 3: In-place operation on a gradient-tracked tensor ---
a = torch.randn(2, 2, requires_grad=True)

# WRONG: in-place add destroys the value autograd needs for gradient computation
# a.add_(torch.ones(2, 2))  # RuntimeError: in-place modification of leaf variable

# CORRECT: create a new tensor — preserves the computation graph
b = a + torch.ones(2, 2)  # new tensor, a unchanged, graph intact
loss = b.sum()
loss.backward()  # gradients computed correctly
print(f"Gradient after correct operation: {a.grad}")

# --- Mistake 4: View vs clone confusion ---
original = torch.tensor([1.0, 2.0, 3.0, 4.0])
view     = original[:2]   # view — shares storage
copied   = original[:2].clone()  # independent copy

view[0] = 99.0   # modifies original too
print(f"After modifying view — original: {original}")  # [99., 2., 3., 4.]

copied[0] = 77.0  # does NOT modify original
print(f"After modifying clone — original: {original}")  # unchanged

# --- Mistake 5: Contiguity and .view() ---
t = torch.randn(3, 4)
t_transposed = t.T   # transpose — non-contiguous, strides reordered
# t_transposed.view(12)  # RuntimeError: view size not compatible with non-contiguous tensor
t_contiguous = t_transposed.contiguous()  # copies data into contiguous layout
reshaped = t_contiguous.view(12)          # works correctly
print(f"Contiguous reshape: {reshaped.shape}")

Control Memory, Control Your Model: The Real Cost of Tensor Shapes

Your model doesn't crash because of a bug. It crashes because you ran out of VRAM at batch 47. Every senior engineer has been there. The fix isn't buying more GPUs; it's understanding how tensor shapes wreck your memory budget. A single (1024, 1024) float32 tensor costs 4 MB. Blow that up to (1024, 2048) and you're at 8 MB. That's fine. But chain a few of these in a transformer and suddenly you're holding 4 GB of intermediate activations. PyTorch's torch.cuda.max_memory_allocated() is your best friend. Call it after every forward pass during development. Watch for the silent killer: broadcasting. A (64, 512) matrix multiplied with (512, 1) creates an implicit (64, 512) output. Every dim mismatch multiplies memory by the batch size. Profile before you optimize. Guesswork is for people who enjoy swapping GPUs out of racks at 3 AM.

// io.thecodeforge — ml-ai tutorial

import torch
import gc

def profile_tensor_memory():
    # Reset the peak memory tracker—don't let previous allocations lie to you
    torch.cuda.reset_peak_memory_stats()
    gc.collect()  # Force Python garbage collector before measurement

    # A typical hidden state in a transformer layer
    batch_size, seq_len, hidden_dim = 64, 128, 1024
    input_tensor = torch.randn(batch_size, seq_len, hidden_dim, device='cuda')

    # This line does a broadcasted addition: waste of memory waiting to happen
    # bias shape (1, 1, hidden_dim) broadcasts to full (64, 128, 1024)
    bias = torch.randn(1, 1, hidden_dim, device='cuda')
    result = input_tensor + bias

    # Peak memory in MB after the operation
    peak_mb = torch.cuda.max_memory_allocated() / (1024 * 1024)
    print(f"Peak GPU memory after broadcasted add: {peak_mb:.2f} MB")

    # Now the same operation with no broadcast—uses the same memory but cleaner
    torch.cuda.reset_peak_memory_stats()
    gc.collect()
    result_direct = input_tensor + bias.expand_as(input_tensor)  # Pre-expand is explicit
    peak_mb_explicit = torch.cuda.max_memory_allocated() / (1024 * 1024)
    print(f"Peak GPU memory with explicit expand: {peak_mb_explicit:.2f} MB")

profile_tensor_memory()

Pin Memory or Pay the Price: The Hidden Cost of CPU-to-GPU Transfer

You think you wrote a fast DataLoader. It uses 8 workers, prefetches 2 batches. But your GPU idle time is 20%. That's because transfers from CPU RAM to GPU VRAM are synchronous by default. Every to(device='cuda') call stalls the GPU until the CPU finishes copying. The fix is pinning memory. When you set pin_memory=True in your DataLoader, PyTorch allocates page-locked memory on the host. That memory is directly accessible by the GPU DMA engine. No page faults, no copy-through overhead. The transfer becomes asynchronous. Your GPU can keep running while the next batch is being prepared. Benchmark this: a ResNet-50 training loop with pin_memory=False vs True. The difference on a 4-GPU node is often 15-25% throughput. Don't let your DataLoader steal cycles from your backprop.

// io.thecodeforge — ml-ai tutorial

import torch
from torch.utils.data import DataLoader, TensorDataset
import time

# Simulate a dataset: 10,000 samples, each 3x224x224 image (standard ImageNet size)
dummy_data = torch.randn(10000, 3, 224, 224)
dummy_labels = torch.randint(0, 1000, (10000,))
dataset = TensorDataset(dummy_data, dummy_labels)

def time_dataloader(pin_memory, num_workers=4):
    loader = DataLoader(
        dataset,
        batch_size=64,
        num_workers=num_workers,
        pin_memory=pin_memory,
        shuffle=True
    )

    start = time.perf_counter()
    for batch_idx, (data, targets) in enumerate(loader):
        # Simulate the transfer to GPU that happens in every training loop
        data = data.to('cuda', non_blocking=True) if pin_memory else data.to('cuda')
        targets = targets.to('cuda', non_blocking=True) if pin_memory else targets.to('cuda')
        # Don't actually train—just time the transfer
        if batch_idx > 50:  # Only sample 50 batches
            break
    end = time.perf_counter()
    return end - start

print("Pinned memory disabled:")
time_no_pin = time_dataloader(pin_memory=False)
print(f"Time for 50 batches: {time_no_pin:.2f}s")

print("\nPinned memory enabled:")
time_pin = time_dataloader(pin_memory=True)
print(f"Time for 50 batches: {time_pin:.2f}s")

print(f"\nSpeedup: {(time_no_pin / time_pin):.2f}x")

Stop Guessing: Profile Tensor Memory Before You Deploy

Memory leaks in production ML don't crash your training job — they crash your inference API at 3 AM when traffic spikes. Most engineers waste days debugging OOM errors that could be caught with one line of instrumentation.

PyTorch's built-in memory profiler (torch.cuda.memory_summary()) shows you exactly where every byte goes. Run it after your model's forward pass. Look for tensors that persist when they shouldn't — those are your memory anchors. The usual suspects: gradients held for backprop when you're in eval mode, or intermediate activations cached by autograd despite torch.no_grad().

Don't trust nvidia-smi. That shows total GPU allocation, not per-tensor breakdown. Use the PyTorch profiler to see allocation by operation. Then kill the hidden tensors. Your production budget will thank you.

// io.thecodeforge — ml-ai tutorial

import torch
import gc

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Simulate leak: hold reference to intermediate tensor
model = torch.nn.Linear(1024, 1024).to(device)
inputs = torch.randn(256, 1024, device=device)

gc.collect()
torch.cuda.empty_cache()

# One forward pass with memory tracking
output = model(inputs)
hidden = output.relu()  # This tensor persists; kill it

del hidden  # Explicitly free
print(torch.cuda.memory_summary(device=device, abbreviated=True))

The One Transform That Breaks Your Batch Norm (And How to Fix It)

Batch normalization tracks running mean and variance per channel. When you reshape a tensor from (N, C, H, W) to (N*H, W, C) for a sequence model, you corrupt those statistics. The channel axis gets scrambled. Your model trains fine but serves garbage.

The fix: never reshape across the channel dimension. If you must flatten spatial dimensions, permute first to move channels to the last axis. Then reshape — the channel content stays contiguous. Or switch to LayerNorm, which normalizes over feature dimensions and doesn't care about spatial layout.

Check your running stats before and after a reshape. If the mean vector changes shape or magnitude, you've got a silent bug. Trust your profiler, not your intuition.

// io.thecodeforge — ml-ai tutorial

import torch
import torch.nn as nn

class BrokenBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.bn = nn.BatchNorm2d(channels)

    def forward(self, x):
        # BAD: reshape collapses channel dim (C) into spatial (H*W)
        b, c, h, w = x.shape
        return self.bn(x).view(b * h, w, c)  # statistic corruption

class FixedBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.bn = nn.BatchNorm2d(channels)

    def forward(self, x):
        b, c, h, w = x.shape
        # Permute channels to last, then reshape spatial without touching C
        x = self.bn(x)
        x = x.permute(0, 2, 3, 1)  # (N, H, W, C)
        return x.reshape(b * h, w, c)

x = torch.randn(2, 4, 8, 16)
broken = BrokenBlock(4)
fixed = FixedBlock(4)
print('Broken output shape:', broken(x).shape)  # (16, 16, 4)
print('Fixed output shape:', fixed(x).shape)    # (16, 16, 4)
# Running means differ — broken one is garbage

Installation: Why the Wrong Build Costs You 10x Latency

Installing PyTorch seems trivial — pip install torch — but the default binary wastes GPU memory and cripples inference speed. The critical decision is selecting the CUDA version that matches your driver and hardware. Use nvidia-smi to check driver-capable CUDA version, then install the corresponding PyTorch build from pytorch.org. On CPU-only systems, avoid the CUDA build entirely; it pulls unnecessary GPU libraries. For edge devices, compile from source with USE_CUDA=0 to cut binary size by 80%. Always verify with torch.cuda.is_available() and torch.backends.cudnn.version(). A mismatch here silently falls back to CPU, multiplying training time. The why: PyTorch is a C++ engine — the Python wheel is just a wrapper. Wrong ABI compatibility forces CPU emulation or crashes. Test on a small tensor: torch.randn(3,3).cuda() should return instantly. If it hangs or errors, your installation is broken.

// io.thecodeforge — ml-ai tutorial

import torch

# Check CUDA availability
print(f"CUDA available: {torch.cuda.is_available()}")

# Get device count
print(f"GPU count: {torch.cuda.device_count()}")

# Quick sanity: create tensor on GPU
try:
    t = torch.randn(4, 4).cuda()
    print("GPU tensor created successfully")
    print(f"Device: {t.device}")
except RuntimeError as e:
    print(f"GPU failure: {e}")

# Display CUDA version used by PyTorch
print(f"PyTorch CUDA version: {torch.version.cuda}")

Enhancing Data Diversity through Augmentation: Why Random Noise Beats Fixed Pipelines

Data augmentation is not about random flips — it's about forcing your model to learn invariances that generalize. Static augmentation pipelines (e.g., always rotate 30°) create spurious correlations. Instead, use stochastic augmentation with per-sample randomness controlled by torch.Generator. The why: Deterministic transforms let the model memorize augmentations as features. Random seeds per batch break that pattern. For images, combine geometric (random affine, perspective) with photometric (color jitter, Gaussian noise) transforms. Always apply augmentations on the CPU with num_workers>0 to avoid blocking GPU compute. Critical: never augment validation or test sets — only training. Use torchvision.transforms.RandAugment for production: it wraps 14 transforms with learned magnitudes. Profile memory: in-place augmentation via torchvision.transforms.functional avoids creating intermediate tensors. The hidden cost: excessive augmentation with RandomResizedCrop can double data loading time if num_workers is under 4.

// io.thecodeforge — ml-ai tutorial

import torch
from torchvision import transforms

# Use a Generator for reproducible randomness per batch
g = torch.Generator().manual_seed(42)

# Stochastic augmentation pipeline (training only)
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5, generator=g),
    transforms.RandAugment(num_ops=2, magnitude=9),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Apply to a batch: shape (B, C, H, W)
# batch_tensors = torch.stack([train_transform(img) for img in batch])