pytorch / pytorch/pytorch

[dynamo] LazyBatchNorm2d fails under torch.compile(dynamic=True) — "SymIntArrayRef expected to contain only concrete integers"

Open
#185,494 3 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: crash module: dynamic shapes module: dynamo module: nn oncall: pt2 triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

LazyBatchNorm2d crashes under torch.compile(dynamic=True) when it has not been pre-initialized. During lazy module initialization, self.num_features becomes a symbolic integer, and torch.empty((self.num_features,)) fails because it requires concrete integers.

Eager mode (without torch.compile ) works correctly.

### Minimal repro
```
import torch, torch.nn as nn

class M(nn.Module):
def __init__(self):
super().__init__()
self.bn = nn.LazyBatchNorm2d()
def forward(self, x):
return self.bn(x)

m = torch.compile(M(), dynamic=True)
m(torch.randn(1, 3, 16, 16)) # !crash here
```
### Error logs
```
Traceback (most recent call last):
...
File "torch/_dynamo/variables/nn_module.py", line 118, in initialize_lazy_module
mod._infer_parameters(mod, fake_args, fake_kwargs)
File "torch/nn/modules/lazy.py", line 263, in _infer_parameters
module.initialize_parameters(*args, **kwargs)
File "torch/nn/modules/batchnorm.py", line 289, in initialize_parameters
self.weight.materialize((self.num_features,))
File "torch/nn/parameter.py", line 147, in materialize
self.data = torch.empty(shape, device=device, dtype=dtype)
torch._dynamo.exc.InternalTorchDynamoError: RuntimeError: /__w/pytorch/pytorch/build/aten/src/ATen/RegisterCPU_1.cpp:2519: SymIntArrayRef expected to contain only concrete integers

from user code:
File "reproduce6.py", line 7, in forward
return self.bn(x)
```
### Root cause
In `batchnorm.py:initialize_parameters` :

```
def initialize_parameters(self, input):
...
self.weight.materialize((self.num_features,)) # ← self.num_features is symbolic
```

When `dynamic=True` , Dynamo traces with fake/symbolic tensors. self.num_features is inferred from the symbolic input shape and remains a SymInt . torch.empty() receives it and triggers the internal assertion.

### Expected behavior
LazyBatchNorm2d should initialize correctly under torch.compile(dynamic=True) , or Dynamo should eagerly resolve num_features to a concrete integer before calling initialize_parameters .

### Versions
```
PyTorch version: 2.13.0.dev20260521+cu126
OS: Linux (Ubuntu 20.04)
Python: 3.12
CUDA: 12.6
```

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @chauhang @penguinwu @ezyang @bobrenjc93 @aditvenk @laithsakka @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @amjames @jataylo @azahed98

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.