PCPJ:

But there is still one thing without answer. Why when I changed the shape of the tensor the Nans disappeared?


Depending on the shapes, different methods will be called as seen here 1.
In particular you could compare the methods used in nll_loss_out_frame 2 and nll_loss2d_forward_out_frame, which might have a different order of accumulation (and the latter seems to over-/underflow in float16) and is why they promote to float32 using amp.



 PCPJ:

Where should I put the with torch.cuda.amp.autocast()?


You can wrap the model execution in it as described in the examples 5:
# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Runs the forward pass with autocasting.
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
        # Backward passes under autocast are not recommended.
        # Backward ops run in the same dtype autocast chose for corresponding forward ops.
        scaler.scale(loss).backward()

        # scaler.step() first unscales the gradients of the optimizer's assigned params.
        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()




 PCPJ:

So, to summarize, I shouldn’t train my model with half precision this way:


You could try to manually cast the model to float16, but as the model can easily create invalid outputs, you would also have to verify, that all operations are working correctly.