code: def test_batch(norm):
      x = torch.randn(3, 8)
      edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
      batch = torch.tensor([0, 0, 1])
  
      model = GraphSAGE(8, 16, num_layers=2, norm=norm)
      assert model.supports_norm_batch == (norm != "BatchNorm")
  
      out = model(x, edge_index, batch=batch)
      assert out.size() == (3, 16)
  
      if model.supports_norm_batch:
          with pytest.raises(RuntimeError, match="out of bounds"):
              model(x, edge_index, batch=batch, batch_size=1)
  
imports: import os
  import os.path
  import random
  import sys
  import warnings
  import pytest
  import torch
  import torch.nn.functional
  import torch_geometric.typing
  from torch_geometric.data import Data
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import SAGEConv
  from torch_geometric.nn.models import GAT, GCN, GIN, PNA, EdgeCNN, GraphSAGE
  from torch_geometric.profile import benchmark
  from torch_geometric.testing import onlyFullTest, onlyLinux, onlyNeighborSampler, onlyOnline, withCUDA, withPackage
label_desc: Use image node_id 9 for calling a global function with example usage: test_batch(norm) without return types