def test_predict_dataset(tmpdir):
    """Tests that we can generate embeddings from a pytorch geometric dataset."""
    tudataset = datasets.TUDataset(root=tmpdir, name="KKI")
    model = GraphEmbedder(
        GraphClassifier(num_features=tudataset.num_features, num_classes=tudataset.num_classes).backbone
    )
    datamodule = DataModule(
        predict_input=GraphClassificationDatasetInput(RunningStage.PREDICTING, tudataset),
        transform=GraphClassificationInputTransform,
        batch_size=4,
    )
    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
    out = trainer.predict(model, datamodule=datamodule)
    assert isinstance(out[0][0], Tensor)
Tests that we can generate embeddings from a pytorch geometric dataset.
tests/graph/embedding/test_model.py