code: def test_lightning_data_custom_sampler():
      class DummySampler(BaseSampler):
          def sample_from_edges(self, *args, **kwargs):
              pass
  
          def sample_from_nodes(self, *args, **kwargs):
              pass
  
      data = Data(
          num_nodes=2, edge_index=torch.tensor([[0, 1], [1, 0]])
      )
  
      datamodule = LightningNodeData(
          data,
          node_sampler=DummySampler(),
          input_train_nodes=torch.arange(2),
      )
      assert isinstance(datamodule.graph_sampler, DummySampler)
  
      datamodule = LightningLinkData(
          data,
          link_sampler=DummySampler(),
          input_train_edges=torch.tensor([[0, 1], [0, 1]]),
      )
      assert isinstance(datamodule.graph_sampler, DummySampler)
  
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from torch import Tensor
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.data.lightning import LightningDataset, LightningLinkData, LightningNodeData
  from torch_geometric.nn import global_mean_pool
  from torch_geometric.sampler import BaseSampler, NeighborSampler
  from torch_geometric.testing import MyFeatureStore, MyGraphStore, get_random_edge_index, onlyCUDA, onlyFullTest, onlyNeighborSampler, onlyOnline, withPackage
label_desc: Use image node_id 4 for calling a global function with example usage: test_lightning_data_custom_sampler() without return types