def METHOD_NAME(model, edges=False, **kwargs):
    dataset = TestDataset(
        [
            _get_graph(
                n_nodes=n_nodes,
                n_features=n_node_features,
                n_edge_features=n_edge_features if edges else None,
            )
            for _ in range(batch_size)
        ]
    )
    loader = loaders.BatchLoader(dataset, epochs=1, batch_size=batch_size)
    inputs = loader.__next__()
    model_instance = model(**kwargs)
    output = model_instance(inputs)