def forward(
        model,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        """Forward pass of the MedusaModel.
        Returns:
            torch.Tensor: A tensor containing predictions from all Medusa heads.
            (Optional) Original predictions from the base model's LM head.
        """
        loss = 0
        medusa_logits = None
        # LOG.debug("medusa_return: %s", medusa_return)
        if not medusa_return:
            return model.old_forward(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        # Pass input through the base model
        if medusa_only_heads:
            with torch.no_grad():
                outputs = model.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_values=past_key_values,
                    inputs_embeds=inputs_embeds,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    return_dict=return_dict,
                )
                hidden_states = outputs[0]
                # The lm_head will be frozen as well, so it's within the context of torch.no_grad()
                if not with_liger:
                    medusa_logits = [model.lm_head(hidden_states)]
        else:
            outputs = model.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
            hidden_states = outputs[0]
            if not with_liger:
                medusa_logits = [model.lm_head(hidden_states)]

        if not with_liger:
            for i in range(model.medusa_num_heads):
                medusa_logits.append(model.medusa_head[i](hidden_states))
            medusa_logits = torch.stack(medusa_logits, dim=0)

        if model.training:
            # Fix all the coefficients to 1 for now
            medusa_scheduler_coefficient = 1
            medusa_heads_coefficient = 1
            medusa_decay_coefficient = 1
            loss = 0

            if with_liger:
                lce = LigerFusedLinearCrossEntropyLoss()
                for i in range(model.medusa_num_heads + 1):
                    shift_hidden_states = (
                        hidden_states[..., : -(1 + i), :].contiguous().view(-1, model.config.hidden_size)
                    )
                    shift_labels = labels[..., (1 + i) :].contiguous().view(-1)

                    weight = model.lm_head.weight if i == 0 else model.medusa_head[i - 1][-1].weight
                    loss_i = lce(weight, shift_hidden_states, shift_labels)

                    loss += calculate_loss_contribution(
                        loss_i,
                        i,
                        medusa_only_heads,
                        medusa_decay_coefficient,
                        medusa_heads_coefficient,
                        medusa_scheduler_coefficient,
                    )
            else:
                loss_fct = CrossEntropyLoss()
                for i in range(model.medusa_num_heads + 1):
                    medusa_logits_i = medusa_logits[i, :, : -(1 + i)].contiguous().view(-1, medusa_logits.shape[-1])
                    medusa_logits_i = medusa_logits_i.float()
                    medusa_labels = labels[..., (1 + i) :].contiguous().view(-1).to(medusa_logits_i.device)

                    loss_i = loss_fct(medusa_logits_i, medusa_labels)

                    loss += calculate_loss_contribution(
                        loss_i,
                        i,
                        medusa_only_heads,
                        medusa_decay_coefficient,
                        medusa_heads_coefficient,
                        medusa_scheduler_coefficient,
                    )
        else:
            if model.config.pretraining_tp > 1:
                raise NotImplementedError
            else:
                medusa_logits = [model.lm_head(hidden_states)]
                for i in range(model.medusa_num_heads):
                    medusa_logits.append(model.medusa_head[i](hidden_states))

        return_dict = return_dict if return_dict is not None else model.config.use_return_dict

        if not return_dict:
            output = (medusa_logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=medusa_logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
Forward pass of the MedusaModel.
        Returns:
            torch.Tensor: A tensor containing predictions from all Medusa heads.
            (Optional) Original predictions from the base model's LM head.
        
examples/medusa/medusa_util.py