code: def __init__(
      self,
      model: torch.nn.Module,
      mask_type: Union[str, MaskLevelType],
      output_idx: Optional[Union[int, Tensor]] = None,
      model_config: Optional[ModelConfig] = None,
  ):
      super().__init__()
  
      self.mask_type = MaskLevelType(mask_type)
      self.model = model
      self.output_idx = output_idx
      self.model_config = model_config
  
imports: from enum import Enum
  from typing import Dict, Optional, Tuple, Union
  import torch
  from torch import Tensor
  from torch_geometric.explain.algorithm.utils import clear_masks, set_hetero_masks, set_masks
  from torch_geometric.explain.config import ModelConfig, ModelMode, ModelReturnType
  from torch_geometric.typing import EdgeType, Metadata, NodeType
label_desc: Use image node_id 1 to create a new CaptumModel object from inherited base classes: torch.nn with example: obj = CaptumModel(model, mask_type, output_idx, model_config)