code: def test_torch_script():
      class Model(torch.nn.Module):
          def forward(self, x: Tensor, edge_index: EdgeIndex) -> Tensor:
              row, col = edge_index[0], edge_index[1]
              x_j = x[row]
              out = scatter(x_j, col, dim_size=edge_index.num_cols)
              return out
  
      x = torch.randn(3, 8)
      # Test that `num_cols` gets picked up by making last node isolated.
      edge_index = EdgeIndex(
          [[0, 1, 1, 2], [1, 0, 0, 1]], sparse_size=(3, 3)
      )
  
      model = Model()
      expected = model(x, edge_index)
      assert expected.size() == (3, 8)
  
      # `torch.jit.script` does not support inheritance at the `Tensor` level :(
      with pytest.raises(
          RuntimeError, match="attribute or method 'num_cols'"
      ):
          torch.jit.script(model)
  
      # A valid workaround is to treat `EdgeIndex` as a regular PyTorch tensor
      # whenever we are in script mode:
      class ScriptableModel(torch.nn.Module):
          def forward(self, x: Tensor, edge_index: EdgeIndex) -> Tensor:
              row, col = edge_index[0], edge_index[1]
              x_j = x[row]
              dim_size: Optional[int] = None
              if not torch.jit.is_scripting() and isinstance(
                  edge_index, EdgeIndex
              ):
                  dim_size = edge_index.num_cols
              out = scatter(x_j, col, dim_size=dim_size)
              return out
  
      script_model = torch.jit.script(ScriptableModel())
      out = script_model(x, edge_index)
      assert out.size() == (2, 8)
      assert torch.allclose(out, expected[:2])
  
imports: import os.path
  import warnings
  from typing import List, Optional
  import pytest
  import torch
  from torch import Tensor, tensor
  import torch_geometric
  from torch_geometric import EdgeIndex
  from torch_geometric.edge_index import SUPPORTED_DTYPES, ReduceType, SortReturnType, _scatter_spmm, _torch_sparse_spmm, _TorchSPMM, set_tuple_item
  from torch_geometric.profile import benchmark
  from torch_geometric.testing import onlyCUDA, onlyLinux, withCUDA, withoutExtensions, withPackage
  from torch_geometric.typing import SparseTensor
  from torch_geometric.utils import scatter
label_desc: Use image node_id 34 for calling a global function with example usage: test_torch_script() and returns: out, out