ML Framework for Varying Batch Size: What Works

The best ML framework for varying batch size is PyTorch: in normal eager mode it runs any batch size with no extra code, and torch.compile handles changing batches with dynamic shapes, either automatically after the first recompile or up front when you mark the batch dimension as dynamic. TensorFlow is a close second, as long as you give tf.function a None batch dimension. JAX and MLX compile a fresh program for each new shape by default, so with them you either pad batches to a few fixed sizes or, in MLX, use shapeless compilation.

So the real question is not which framework can run a batch of 7 after a batch of 32. They all can. The question is what each one does behind the scenes when the size changes: nothing, a silent recompile, or an error. This guide explains that for each framework, with the settings that fix it, based on the official documentation checked on November 4, 2026. For a broader comparison of these libraries, see our guide to the best ML frameworks for Python.

ML framework for varying batch size: quick comparison

Why batch size changes in the first place

There are three common causes, and they need different fixes.

1. The last batch of an epoch. If your dataset size is not divisible by the batch size, the final batch is smaller. PyTorch's DataLoader keeps it by default and drops it only when you set drop_last=True, according to the torch.utils.data documentation.
2. Inference traffic. A server batches whatever requests arrive, so batch sizes go up and down all day.
3. Variable sequence length. In language models, the batch dimension often stays put while the sequence dimension changes. Every fix below applies to any dimension, not only the first one.

If cause 1 is your only problem, drop_last=True during training is the simplest answer and you can stop here. Causes 2 and 3 need the framework to cope.

PyTorch: eager just works, torch.compile needs a hint

In eager mode, PyTorch runs each operation as Python reaches it, so a batch of 1, 7 or 4,096 all work as long as the GPU has the memory. Most people never think about it.

torch.compile is different. The Dynamic Shapes Core Concepts page says PyTorch assumes static shapes by default and, when a size change is detected, tries to recompile with dynamic input. The torch.compile reference spells out the three settings of the dynamic argument:

With the default, you pay for two compilations: one for the first batch size and one more general version after the size changes. That is usually fine. If you know the batch will vary, mark it up front so the first compile is already general:

import torch

model = torch.compile(model)

x = next(batches)  # shape [batch, seq, hidden]
torch._dynamo.mark_dynamic(x, 0)  # batch dimension varies
out = model(x)

The docs warn that dynamic=True is most useful for small models and may cause crashes or performance problems on large ones, so prefer mark_dynamic on the specific dimension. To see why a recompile happened, run with TORCH_LOGS=recompiles or TORCH_LOGS=dynamic, as the recompilation guide describes. There is also a cap on how many times a function can recompile; once it is hit, Dynamo stops compiling that function.

The batch size of 1 trap

PyTorch treats sizes 0 and 1 as special. The zero one specialization page explains that torch.compile automatically specializes on inputs of size 0 or 1 and assumes other sizes are not 0 or 1. In practice, a batch of 1 (a single request, or a tiny last batch) can trigger its own compile even when larger batches share one dynamic graph. If you serve single requests often, warm up with a batch of 1 as well as a larger batch.

CUDA graphs need fixed shapes

mode="reduce-overhead" adds CUDA graphs, which cut Python overhead for small batches. The CUDAGraph Trees documentation says it records a new graph for every unique input shape, because CUDA graphs need fixed memory addresses. Its own advice for frequently changing shapes is to pad inputs to a few fixed sizes. You can also set torch._inductor.config.triton.cudagraph_skip_dynamic_graphs=True so only static shape functions get graphed.

TensorFlow and Keras: give tf.function a None batch dimension

TensorFlow's eager mode accepts any batch size. The catch is tf.function, which turns Python into a graph. The tf.function guide explains that each traced graph is specialized to its input types and shapes, and tracing is expensive enough that retracing on every call can make the code slower than not using tf.function at all.

Two fixes, both from that guide:

@tf.function(input_signature=(
    tf.TensorSpec(shape=[None, 128], dtype=tf.float32),
))
def predict(x):
    return model(x)

A None dimension acts as a wildcard, so any batch size reuses the same graph. If you cannot write a signature, @tf.function(reduce_retracing=True) traces once with the first shape, then traces a more general graph and reuses it after that. The guide notes this is less efficient than setting the signature directly.

Keras 3 runs on JAX, TensorFlow or PyTorch, per the Keras getting started page, so how it handles a changing batch follows the backend you choose. If varying batches matter, the PyTorch or TensorFlow backend is the easier ride. Our TensorFlow vs PyTorch vs Keras guide covers how the three relate.

JAX: one compile per shape, so bucket your batches

JAX is the framework where varying batch size costs the most. The JAX shape polymorphism docs state that in JIT mode a function is traced, lowered and compiled for each combination of input types and shapes. A new batch size means a new compile, and the JIT compilation guide warns that when there are many distinct values, a program can spend more time compiling than running.

Shape polymorphism in jax.export lets you trace and lower once with a symbolic batch dimension, but the same page says the exported function is still recompiled for each concrete shape it meets. It saves tracing, not compiling.

The practical answer is bucketing:

1. Pick a few batch sizes, for example 8, 16, 32 and 64.
2. Pad each real batch up to the nearest bucket with dummy rows.
3. Pass a mask so padded rows do not count in the loss or metrics.
4. Drop the padded rows from the output.

With four buckets, JAX compiles four programs once and reuses them. This is also why JAX shines on steady, large training jobs and feels awkward for bursty inference. Our PyTorch vs JAX comparison goes into that tradeoff.

MLX on a Mac: shapeless compilation

MLX runs eagerly by default, so changing batch sizes is no problem there. mx.compile is the exception: the MLX compile docs list a change in shape or number of dimensions as a reason to recompile. Pass shapeless=True and shape changes no longer trigger a recompile. The docs warn that any code that depends on input shapes, such as a reshape computed from x.shape, will break or give wrong results under shapeless compilation, so use it only on shape agnostic functions.

Inference engines: TensorRT and vLLM

If you export a model for deployment, the runtime has its own rules.

TensorRT needs you to declare the range. The TensorRT dynamic shapes docs say you mark runtime dimensions with -1 and create at least one optimization profile with minimum, optimum and maximum shapes. The builder tunes kernels for the optimum shape, and inputs outside the min to max range fail at runtime. Some kernels only work when min, opt and max are equal, so a wide range can cost speed and memory. A common pattern is several profiles, one per batch range.

vLLM and similar LLM servers handle varying batches for you. Its CUDA Graphs design notes describe dispatching each step to a graph keyed on a padded token count from a set of capture sizes, the same bucketing idea as above, done automatically. If you serve open weight models, you rarely need to touch this. Our vLLM vs llama.cpp guide covers choosing a server.

How to pick bucket sizes

Padding wastes compute on dummy rows, so choose buckets with care:

1. Log the real batch sizes your job sees for a day or an epoch.
2. Put buckets where the traffic is. Powers of two are a common default, but buckets near your most frequent sizes waste less.
3. Keep the count small. Every bucket is one more compile and, with CUDA graphs, one more recorded graph that uses memory.
4. Check memory at the largest bucket. Activation memory grows with batch size, and our guide on how much VRAM an LLM needs shows how to estimate it.

Which should you choose?

• You train or fine tune and want zero fuss: PyTorch in eager mode, adding torch.compile with mark_dynamic once the code works.
• You already run TensorFlow: keep it, and put None in every tf.function input signature.
• You train at scale on TPUs: JAX, with bucketed batches and masks.
• You are on Apple silicon: MLX, using shapeless=True only where the code has no shape logic.
• You deploy a fixed model to NVIDIA GPUs: TensorRT with one or more optimization profiles.

FAQ

Can PyTorch handle a different batch size every step?

Yes. Eager PyTorch accepts any batch size with no changes. With torch.compile, the default setting recompiles once when it first sees a new size and then uses a dynamic kernel. Mark the batch dimension with torch._dynamo.mark_dynamic to skip that extra compile.

Why does my compiled model recompile on a batch size of 1?

torch.compile specializes on sizes 0 and 1, so a batch of 1 can get its own graph even when other sizes share a dynamic one. Warm up with a batch of 1 if you serve single requests, or drop the tiny last batch in training with drop_last=True.

Does JAX support dynamic batch size?

JAX runs any batch size, but jax.jit compiles a separate program for each new input shape. For workloads with many sizes, pad batches to a few fixed bucket sizes and mask the padded rows so you only pay for a handful of compiles.

Should I just drop the last incomplete batch?

During training it is often the simplest fix: set drop_last=True in the PyTorch DataLoader and every batch has the same size. During evaluation or inference you usually need every example, so pad or use dynamic shapes instead.

Does padding change the model's results?

It should not if you mask correctly. Padded rows must be excluded from the loss, metrics and outputs. Layers that compute statistics across the batch, such as batch normalization in training mode, are the main place where padding or a very small batch can shift results.