No no I’m not doing half() anywhere, It means I’m using the procedure from here: Automatic Mixed Precision examples — PyTorch 1.8.0 documentation 2.
My code:
Main while loop:
    fetching batch from dataloader
    batch['x'] = batch['x'].cuda(non_blocking=True)

    #infer
    with torch.cuda.amp.autocast():
        out = model(batch['x'])
        loss = loss_model(out, batch['x']) #autoencoder loss
    
    #backprop
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5) 
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

batch[‘x’] is of shape [64, 1, 256, 150] with values ranging from 0 to 1e4.
I know the troubles with manually using half() on a model