[dynamo] LazyBatchNorm2d fails under torch.compile(dynamic=True) — "SymIntArrayRef expected to contain only concrete integers"
- 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
Assessment
This issue has not been assessed yet.