Function test_execute

tests/python/relax/test_training_trainer_numeric.py:57–77

(target, dev)

@tvm.testing.parametrize_targets("llvm")
def test_execute(target, dev):
    backbone = _get_backbone()
    pred_sinfo = relax.TensorStructInfo((1, 5), "float32")

    setup_trainer = SetupTrainer(
        MSELoss(reduction="sum"),
        Adam(0.01),
        [pred_sinfo, pred_sinfo],
    )

    train_mod = setup_trainer(backbone)
    ex = tvm.compile(train_mod, target)
    vm = relax.VirtualMachine(ex, dev)

    trainer = Trainer(train_mod, vm, dev, False)
    trainer.zero_init_params()
    trainer.xaiver_uniform_init_params()

    dataset = _make_dataset()
    trainer.predict(dataset[0][0])
    trainer.update(dataset[0][0], dataset[0][1])


@tvm.testing.parametrize_targets("llvm")