code: def test_hetero_linear_amp(device, use_segment_matmul):
      warnings.filterwarnings("ignore", ".*but CUDA is not available.*")
  
      old_state = torch_geometric.backend.use_segment_matmul
      torch_geometric.backend.use_segment_matmul = use_segment_matmul
  
      x = torch.randn(3, 16, device=device)
      type_vec = torch.tensor([0, 1, 2], device=device)
  
      lin = HeteroLinear(16, 32, num_types=3).to(device)
  
      with torch.cuda.amp.autocast():
          assert lin(x, type_vec).size() == (3, 32)
  
      torch_geometric.backend.use_segment_matmul = old_state
  
imports: import copy
  import warnings
  from typing import List
  import pytest
  import torch
  from torch import Tensor
  from torch.nn import Linear
  from torch.nn.parameter import UninitializedParameter
  import torch_geometric.backend
  from torch_geometric.nn import HeteroDictLinear, HeteroLinear, Linear
  from torch_geometric.profile import benchmark
  from torch_geometric.testing import withCUDA, withPackage
  from torch_geometric.typing import pyg_lib
  from torch_geometric.utils import cumsum
label_desc: Use image node_id 9 for calling a global function with example usage: test_hetero_linear_amp(device, use_segment_matmul) without return types