def METHOD_NAME(weights, batch, targets, lr=0.2):
    def compute_loss(weights, batch, targets):
        output = functional_call(model, weights, batch)
        loss = loss_fn(output, targets)
        return loss
    grad_weights, loss = grad_and_value(compute_loss)(weights, batch, targets)
    # NB: PyTorch is missing a "functional optimizer API" (possibly coming soon)
    # so we are going to re-implement SGD here.
    new_weights = {}
    with torch.no_grad():
        for key in grad_weights:
            new_weights[key] = weights[key] - grad_weights[key] * lr
    return loss, new_weights