Function func

imperative/python/test/integration/test_xla_dist_training.py:75–81

(model, optimizer, timage, tlabel)

            label = np.random.randint(0, 10, (3, 8,))

            def func(model, optimizer, timage, tlabel):
                with gm:
                    score = model(timage)
                    loss = F.nn.cross_entropy(score, tlabel)
                    gm.backward(loss)
                    optimizer.step().clear_grad()
                return loss

            if is_trace:
                func = xla_trace(func, without_host=True, capture_as_const=True)