code: def test_neighbor_loader_input_id():
      data = HeteroData()
      data["a"].num_nodes = 10
      data["b"].num_nodes = 12
  
      row = torch.randint(0, data["a"].num_nodes, (40,))
      col = torch.randint(0, data["b"].num_nodes, (40,))
      data["a", "b"].edge_index = torch.stack([row, col], dim=0)
      data["b", "a"].edge_index = torch.stack([col, row], dim=0)
  
      mask = torch.ones(data["a"].num_nodes, dtype=torch.bool)
      mask[0] = False
  
      loader = NeighborLoader(
          data,
          input_nodes=("a", mask),
          batch_size=2,
          num_neighbors=[2, 2],
      )
      for i, batch in enumerate(loader):
          if i < 4:
              expected = [(2 * i) + 1, (2 * i) + 2]
          else:
              expected = [(2 * i) + 1]
  
          assert batch["a"].input_id.tolist() == expected
  
imports: import os.path
  import numpy
  import pytest
  import torch
  import torch_geometric.typing
  from torch_geometric import EdgeIndex
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import GraphConv, to_hetero
  from torch_geometric.sampler.base import SubgraphType
  from torch_geometric.testing import MyFeatureStore, MyGraphStore, get_random_edge_index, get_random_tensor_frame, onlyLinux, onlyNeighborSampler, onlyOnline, withCUDA, withPackage
  from torch_geometric.typing import WITH_EDGE_TIME_NEIGHBOR_SAMPLE, WITH_PYG_LIB, WITH_TORCH_SPARSE, WITH_WEIGHTED_NEIGHBOR_SAMPLE, TensorFrame
  from torch_geometric.utils import is_undirected, sort_edge_index, to_torch_csr_tensor, to_undirected
label_desc: Use image node_id 21 for calling a global function with example usage: test_neighbor_loader_input_id() without return types