Mixed precision training #127

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()

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