I am training with half precision because the final hardware I will deploy my model works with half precision.
I was doing this cast manually because I though it was the right way.
Sorry about that.
But there is still one thing without answer. Why when I changed the shape of the tensor the Nans disappeared?
So, to summarize, I shouldn’t train my model with half precision this way:
model = Model()
model.to('cuda').half()
for layer in model.modules():
    if isinstance(layer, torch.nn.BatchNorm2d):
        layer.float()

for i, (x y) in enumerate(trainLoader):
        x = x.to('cuda').half()
        y = y.to('cuda').half()

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