Solved by ptrblck in post #4 
              
              
                No idea at the moment and you could use torch.autograd.set_detect_anomaly(True) at the beginning of the script to get a stack trace, which should hopefully point to the first NaN creation. 

momentum is an argument in all batchnorm layers, which you can directly pass to them: 
nn.BatchNorm2d(num_fe…