d2l-ai/d2l-zh · error · AssertionError

X.shape[0] == y.shape[0]

Error message

X.shape[0] == y.shape[0]

What it means

AssertionError from d2l.torch.split_batch (sec_multi_gpu): X and y must share the same leading batch dimension before nn.parallel.scatter distributes them across devices. The assert is a pre-flight check so mismatched data fails loudly instead of producing silent device-allocation errors inside scatter.

Source

Thrown at d2l/torch.py:1480

class Benchmark:
    """用于测量运行时间"""
    def __init__(self, description='Done'):
        """Defined in :numref:`sec_hybridize`"""
        self.description = description

    def __enter__(self):
        self.timer = d2l.Timer()
        return self

    def __exit__(self, *args):
        print(f'{self.description}: {self.timer.stop():.4f} sec')

def split_batch(X, y, devices):
    """将X和y拆分到多个设备上

    Defined in :numref:`sec_multi_gpu`"""
    assert X.shape[0] == y.shape[0]
    return (nn.parallel.scatter(X, devices),
            nn.parallel.scatter(y, devices))

def resnet18(num_classes, in_channels=1):
    """稍加修改的ResNet-18模型

    Defined in :numref:`sec_multi_gpu_concise`"""
    def resnet_block(in_channels, out_channels, num_residuals,
                     first_block=False):
        blk = []
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.append(d2l.Residual(in_channels, out_channels,
                                        use_1x1conv=True, strides=2))
            else:
                blk.append(d2l.Residual(out_channels, out_channels))
        return nn.Sequential(*blk)

View on GitHub (pinned to e6b18ccea7)

Solutions

  1. Slice both tensors identically: split_batch(X[:n], y[:n], devices)
  2. Always obtain X, y together from one DataLoader batch rather than separate iterators
  3. Validate shapes at the source: if X.shape[0] != y.shape[0], raise a contextual error including shapes and step index
  4. Check intermediate reshapes keep the batch axis aligned (y.view(-1) vs y.view(1, -1))

Example fix

# before
X = X[:32]  # manual trim of features only
split_batch(X, y, devices)  # y.shape[0]=256 -> AssertionError
# after
X, y = X[:32], y[:32]
split_batch(X, y, devices)  # both 32 rows, scatter ok
Defensive patterns

Strategy: validation

Validate before calling

assert X.shape[0] == y.shape[0], f'misaligned batch: X={tuple(X.shape)} y={tuple(y.shape)}'
Xs, ys = split_batch(X, y, devices)

Type guard

def batch_aligned(X, y) -> bool:
    return X.shape[0] == y.shape[0]

Prevention

When it happens

Trigger: split_batch(X[:64], y, devices) after trimming only X; y coming from a different batch than X (double next(iter(...)) calls); labels reshaped to (n,1) while X stays (n, c, h, w) is fine, but y flattened to a scalar or None breaks it; mixing numpy arrays (no .shape ordering guarantee after slicing mistakes) with tensors.

Common situations: Custom multi-GPU training loops modeled on the book; asynchronous prefetch where X is fetched one step ahead of y; last-batch handling that drops unmatched rows from only one tensor; debugging sessions that slice batches manually.

Related errors


AI-assisted analysis of d2l-ai/d2l-zh@e6b18ccea7 (2026-08-14). Data as JSON: /api/errors/4f79464c62503e1f. Report an issue: GitHub.