Function f

imperative/python/test/unit/amp/test_grad_scaler.py:17–31

(idx, data, calc)

    scaler = GradScaler()

    def f(idx, data, calc):
        x = mge.tensor(data, no_cache=True)
        y = mge.tensor(data, no_cache=True)

        if is_trace:
            calc = trace(calc)

        gm.attach([x, y])
        with gm:
            loss = calc(x, y)
            scaler.backward(gm, loss, unscale_grad=False)
        np.testing.assert_equal(x.grad.numpy(), 2 * scaler.scale_factor)
        scaler.unscale(filter(lambda t: t.grad is not None, gm.attached_tensors()))
        # scaler.unscale(gm.attached_tensors())
        np.testing.assert_equal(x.grad.numpy(), 2)

    def double_variables(x, y):
        z = x + 2 * y