I am training a model with mixed precision (amp_level=‘O1’) with a stack of 3D ResNet blocks.

```
class BasicBlock3d(nn.Module):
""" 3x3x3 Resnet Basic Block"""
expansion = 1
__constants__ = ['downsample']
def __init__(self, inplanes, planes, stride=1, downsample=None,dilation=1):
super(BasicBlock3d, self).__init__()
self.conv1 = conv3x3x3(inplanes, planes, stride, 1, dilation)
self.bn1 = nn.SyncBatchNorm(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3x3(planes, planes, 1, 1, dilation)
self.bn2 = nn.SyncBatchNorm( planes)
if downsample is not None:
self.downsample = downsample
elif stride != 1 or inplanes != planes:
self.downsample = nn.Sequential(conv1x1x1(inplanes, planes, stride), get_norm_3d(norm, planes))
else:
self.downsample = None
self.stride = stride
def forward(self, x):
identity = x
out1 = self.conv1(x)
out2 = self.bn1(out1)
out_ = self.relu(out2)
out3 = self.conv2(out_)
out = self.bn2(out3)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
```

After training more than tens of epochs (70+), the `self.bn1.running_var`

turns to be NaN in some middle blocks and crash the training pipeline (loss=NaN). I have found the direct reason is the input feature (`out1`

) of `self.bn1`

got a sum larger than 65,504 in certain channel, which leads its mean becomes `inf`

. One example is shown below

```
In[32]:out1[0,15,0,1,:]
Out[32]:
tensor([14720., 21936., 21776., 21584., 21376., 21120., 20848., 20544., 20272.,
19984., 19760., 19568., 19376., 19168., 18928., 18656., 18416., 18192.,
18032., 17904., 17776., 17632., 17504., 17392., 17296., 17264., 17216.,
17168., 17120., 17088., 17072., 17056., 17040., 17072., 17136., 17216.,
17264., 17408., 17584., 17824., 17968., 18080., 18176., 18240., 18304.,
18352., 18400., 18448., 18528., 18624., 18784., 18992., 19232., 19520.,
19840., 20192., 20560., 20976., 21408., 21888., 22336., 22720., 22960.,
15320.], dtype=torch.float16, grad_fn=<SliceBackward>)
In[33]:torch.mean(out1[0,15,0,1,:])
Out[33]: tensor(inf, dtype=torch.float16, grad_fn=<MeanBackward0>)
In[35]:torch.sum(out1[0,15,0,1,:])
Out[35]: tensor(inf, dtype=torch.float16, grad_fn=<SumBackward0>)
In[36]:torch.sum(out1[0,15,0,1,:].float())
Out[36]: tensor(1206136., grad_fn=<SumBackward0>)
In[38]:out2[0,15].isnan().all()
Out[38]: tensor(True)
```

Is there any way we can prevent the feature growing this large during the training process?