code: def __init__(self, config, weights, lm_head):
      super().__init__()
      self.heads = torch.nn.ModuleList(
          [
              MedusaHead(config, prefix=f"{i}", weights=weights)
              for i in range(config["medusa_num_heads"])
          ]
      )
      self.lm_head = lm_head
  
file_path: text-generation-inference/server/text_generation_server/utils/medusa.py
imports: import torch
  from dataclasses import dataclass
  from text_generation_server.utils.layers import TensorParallelHead, FastLinear
label_desc: Use image node_id 1 to create a new MedusaModel object from inherited base classes: torch.nn with example: obj = MedusaModel(config, weights, lm_head)