I am getting AssertionError: No inf checks were recorded for this optimizer. in "/torch/cuda/amp/grad_scaler.py", line 291 when mixed-precision is used in this weird example below. However, if no mixed-precision is used pytorch doesn’t complain (toggle USE_HALF_PRECISION = True).
I am using PyTorch 1.6.0 (python 3.7, cuda 10.2.89, cudnn 7.6.5. – everything is in conda binaries). Here is the MWE.
import torch
from torch import nn
from torch.cuda.amp.autocast_mode import autocast
from torch.cuda.amp.grad_scaler import GradScaler

class Identity_with_weights(nn.Module):
    '''For example a KNN algorithm which returns a closest entry from a database for x. Weights are needed
    for a seamless inclusion of knn baseline to a set of baseline which do have some parameters. Otherwise
    you would need to change the code (remove optimizer, backward pass etc) just for knn which is not neat.'''
    def __init__(self):
        super(Identity_with_weights, self).__init__()
        self.__hidden__ = torch.nn.Linear(1, 1, bias=False)

    def forward(self, x):
        # we need it to be able to call backward on the loss which uses x (outputs).
        # Nothing will happen in these examples as it propagates to the input which is not used anywhere else
        x.requires_grad = True
        return x


if __name__ == "__main__":
    # config
    USE_HALF_PRECISION = True
    device = torch.device('cuda:0')

    # define model
    model = Identity_with_weights()

    # define training things
    criterion = nn.L1Loss()
    optimizer = torch.optim.Adam(model.parameters())

    # for amp
    scaler = GradScaler()

    # targets are exactly the same as inputs, i.e. for reconstruction
    inputs = torch.rand(8, 1)
    targets = inputs.clone().detach()

    # send to device
    model = model.to(device)
    inputs = inputs.to(device)
    targets = targets.to(device)

    # we don't need it for the sake of this example, but let's have it here anyway.
    optimizer.zero_grad()

    # since outputs are going to be f16 and targets are f32, criterion will output non zero loss
    if USE_HALF_PRECISION:
        targets = targets.half()
        inputs = inputs.half()

    # autocasting ops inside of the context manager
    with autocast(USE_HALF_PRECISION):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    print(loss)

    # scaling loss if using half precision
    if USE_HALF_PRECISION:
        scaler.scale(loss).backward()
        scaler.step(optimizer) ## ERROR HERE
        scaler.update()
    else:
        loss.backward()
        optimizer.step()

I think I am doing something wrong here. What does it complain about?