code: def __init__(
      self,
      module: torch.nn.Module,
      types: Union[List[NodeType], List[EdgeType]],
  ):
      from torch_geometric.nn import HeteroLinear, Linear
  
      super().__init__()
  
      self.types = types
  
      if isinstance(module, Linear):
          in_channels = module.in_channels
          out_channels = module.out_channels
          bias = module.bias is not None
  
      elif isinstance(module, torch.nn.Linear):
          in_channels = module.in_features
          out_channels = module.out_features
          bias = module.bias is not None
  
      else:
          raise ValueError(
              f"Expected 'Linear' module (got '{type(module)}'"
          )
  
      # TODO We currently assume that `x` is sorted according to `type`.
      self.hetero_module = HeteroLinear(
          in_channels,
          out_channels,
          num_types=len(types),
          is_sorted=True,
          bias=bias,
      )
  
imports: import copy
  import warnings
  from typing import Dict, List, Optional, Union
  import torch
  import torch.nn.functional
  from torch import Tensor
  import torch_geometric
  from torch_geometric import is_compiling
  from torch_geometric.typing import EdgeType, NodeType, OptTensor
  from torch_geometric.utils import cumsum, scatter
label_desc: Use image node_id 1 to create a new ToHeteroLinear object from inherited base classes: torch.nn with example: obj = ToHeteroLinear(module, types)