Function trace_func

imperative/python/test/unit/xla/test_xla_trace.py:152–157

(model, inp1, inp2, scale_factor)

        @xla_trace(without_host=True)
        def trace_func(model, inp1, inp2, scale_factor):
            gm.attach([inp1, inp2])
            with gm:
                loss = model(inp1, inp2)
                scaler.backward(gm, scale_factor, loss)
            return [inp1.grad, scale_factor]

        mge_rsts = trace_func(test_module, x, y, scaler.scale_factor)
        xla_rsts = trace_func(test_module, x, y, scaler.scale_factor)