code: def test_predict_dataset(tmpdir):
      """Tests that we can generate predictions from a pytorch geometric dataset."""
      tudataset = datasets.TUDataset(root=tmpdir, name="KKI")
      model = GraphClassifier(
          num_features=tudataset.num_features,
          num_classes=tudataset.num_classes,
      )
      datamodule = DataModule(
          predict_input=GraphClassificationDatasetInput(
              RunningStage.TESTING, tudataset
          ),
          transform=GraphClassificationInputTransform,
          batch_size=4,
      )
      trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
      out = trainer.predict(
          model, datamodule=datamodule, output="classes"
      )
      assert isinstance(out[0][0], int)
  
file_path: lightning-flash/tests/graph/classification/test_model.py
imports: from typing import Any
  import pytest
  import torch
  from flash import RunningStage, Trainer
  from flash.core.data.data_module import DataModule
  from flash.core.data.io.input import DataKeys
  from flash.core.utilities.imports import _TOPIC_GRAPH_AVAILABLE
  from flash.graph.classification import GraphClassifier
  from flash.graph.classification.input import GraphClassificationDatasetInput
  from flash.graph.classification.input_transform import GraphClassificationInputTransform
  from torch import Tensor
  from tests.helpers.task_tester import TaskTester
label_desc: Use image node_id 1 for calling a global function with example usage: test_predict_dataset(tmpdir) without return types