why does batch normalization make my batches so abnormal?
I'm playing with pytorch for the first time, and I've noticed that when training my neural net, about one time in four or so the loss takes a left turn towards infinity, then nan shortly after that. I've seen a few other questions about nan-ing, but the recommendations there seem to be essentially to do normalization; but the first layer in my net below is such a normalization, and I still see this problem! The full net is a bit convoluted, but I've done some debugging to try to produce a very small, understandable net that still displays the same issue.
The code is below; it consists of sixteen inputs, 0-1, which are passed through a batch normalization and then a fully-connected layer to a single output. I'd like it to learn the function that always outputs 1, so I take the squared error from 1 for the loss.
import torch as t
import torch.nn as tn
import torch.optim as to

if __name__ == '__main__':
    board = t.rand([1,1,1,16])
    net = tn.Sequential \
        ( tn.BatchNorm2d(1)
        , tn.Conv2d(1, 1, [1,16])
        )
    optimizer = to.SGD(net.parameters(), lr=0.1)
    for i in range(10):
        net.zero_grad()
        nn_outputs = net.forward(board)
        loss = t.sum((nn_outputs - 1)**2)
        print(i, nn_outputs, loss)
        loss.backward()
        optimizer.step()

If you run it a few times, eventually you'll see a run that looks like this:
0 tensor([[[[-0.7594]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(3.0953, grad_fn=&lt;SumBackward0&gt;)
1 tensor([[[[4.0954]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(9.5812, grad_fn=&lt;SumBackward0&gt;)
2 tensor([[[[5.5210]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(20.4391, grad_fn=&lt;SumBackward0&gt;)
3 tensor([[[[-3.4042]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(19.3966, grad_fn=&lt;SumBackward0&gt;)
4 tensor([[[[823.6523]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(676756.7500, grad_fn=&lt;SumBackward0&gt;)
5 tensor([[[[3.5471e+08]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(1.2582e+17, grad_fn=&lt;SumBackward0&gt;)
6 tensor([[[[2.8560e+25]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(inf, grad_fn=&lt;SumBackward0&gt;)
7 tensor([[[[inf]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(inf, grad_fn=&lt;SumBackward0&gt;)
8 tensor([[[[nan]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(nan, grad_fn=&lt;SumBackward0&gt;)
9 tensor([[[[nan]]]], grad_fn=&lt;MkldnnConvolutionBackward&gt;) tensor(nan, grad_fn=&lt;SumBackward0&gt;)

Why does my loss go to nan, and what can I do about it?

Welcome to pytorch!
Here is how I would set up your training. Please check the comments.
# how the comunity usually does the import:
import torch  # some people do: import torch as th 
import torch.nn as nn
import torch.optim as optim

if __name__ == '__main__':
    
    # setting some parameters:
    batch_size = 32
    n_dims = 128

    # select GPU if available
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # initializing a simple neural net
    net = nn.Sequential(nn.Linear(n_dims, n_dims  // 2), # Batch norm is not usually used directly on the input
                        nn.BachNorm1d(n_dims  // 2), # Batch norm is used before the activation function (it centers the input and helps make the dims of the previous layers independent of each other)
                        nn.ReLU(), # the most common activation function
                        nn.Linear(n_dims  // 2, 1)  # final layer)
    net.to(device) # model is copied to the GPU if it is availalbe

    optimizer = to.SGD(net.parameters(), lr=0.01) # it is better to start with a low lr and increase it at later experiments to avoid training divergence, the range [1.e-6, 5.e-2] is recommended.

    for i in range(10): 
       # generating random data:
       board = torch.rand([batch_size, n_dims])
       # for sequences: [batch_size, channels, L]
       # for image data: [batch_size, channels, W, H]
       # for videos: [batch_size, chanels, L, W, H]
       boad = board.to(device) # data is copied to the gpu if it is available

        optimizer.zero_grad()  # the convension the comunity uses, though the result is the same as net.zero_grad() 

        nn_outputs = net(board) # don't call net.forward(x), call net(x). Pytorch applies some hooks in the net.__call__(x) that are useful for backpropagation.

        loss = ((nn_outputs - 1)**2).mean() # using .mean() makes your training less sensitive to the batch size.

        print(i, nn_outputs, loss.item())

        loss.backward()
        optimizer.step()

One comment about the batch norm. Per dimension, it calculates the mean and the standard deviation of your batch (check the documentation https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html#torch.nn.BatchNorm2d):
x_normalized = (x.mean(dim=0) / (x.std(dim=0) + e-6)) * scale + shift

Where scale and shift are learnable parameters. If you only give one example per batch, x.std(0) = 0 will make x_normalized contain very very large values.
