code: def forward(self, x):
      logits = self.lm_head(x)
      speculative_logits = torch.stack(
          [head(x) for head in self.heads], dim=1
      )
      return logits, speculative_logits
  
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 2 for calling the MedusaModel obj's underlying member method code with example usage: obj.forward(x) and returns: logits,  speculative_logits