Function train_func

imperative/python/test/integration/test_dtr.py:111–117

(data, label, *, net, gm)

    gm = GradManager().attach(resnet1202.parameters())

    def train_func(data, label, *, net, gm):
        net.train()
        with gm:
            pred = net(data)
            loss = F.loss.cross_entropy(pred, label)
            gm.backward(loss)
        return pred, loss

    _, free_mem = mge.device.get_mem_status_bytes()
    tensor_mem = free_mem - (2 ** 30)