Mixed precision training

[zhiqihuang]

I want to use torch mixed precision training with higher. Is there any example code snippet I can follow?

btw, I run code like this:

scaler = torch.cuda.amp.GradScaler()

with higher.innerloop_ctx(model, optim) as (fmodel, diffopt):
    with autocast():
        logits = fmodel(data)
        loss = loss_function(logits, ys)
    scaler.scale(loss).backward()
    scaler.unscale_(diffopt)
    scaler.step(diffopt)
    scaler.update()
    diffopt.zero_grad()

And get the error message

if param.grad is None:
    AttributeError: 'NoneType' object has no attribute 'grad'

at step scaler.unscale_(diffopt)

Thanks for the help.