Function _export_and_execute

runtime/test/test_runtime_xnnpack.py:28–49

Export *model* with XNNPACK, save to a temp .pte, and run it via Runtime.

(
    model: torch.nn.Module, example_inputs: tuple[torch.Tensor, ...]
)

def _export_and_execute(
    model: torch.nn.Module, example_inputs: tuple[torch.Tensor, ...]
):
    """Export *model* with XNNPACK, save to a temp .pte, and run it via Runtime."""
    with tempfile.TemporaryDirectory() as temp_dir, torch.no_grad():
        model.eval()
        aten = export(model, example_inputs, strict=True)
        edge = to_edge_transform_and_lower(
            aten,
            compile_config=EdgeCompileConfig(_check_ir_validity=False),
            partitioner=[XnnpackPartitioner()],
        )
        et = edge.to_executorch()

        pte_path = Path(temp_dir) / "xnnpack_runtime_test.pte"
        et.save(str(pte_path))

        runtime = Runtime.get()
        program = runtime.load_program(pte_path, verification=Verification.Minimal)
        method = program.load_method("forward")
        assert method is not None, "forward method should exist in exported program"
        return method.execute(example_inputs)


@unittest.skipUnless(