diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 2053ddbc4..7ef5e282e 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -821,6 +821,9 @@ class MiniMaxM2ForCausalLM(nn.Module): else: self.model.layers_to_capture = [val + 1 for val in layer_ids] + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + @torch.no_grad() def forward( self,