fix jamba accuracy test failure (#3466)

Author: blzheng

examples/cpu/llm/inference/distributed/run_accuracy_with_deepspeed.py

@@ -730,6 +730,8 @@ def _model_call(
             example_dict["output_router_logits"] = torch.tensor(
                 model_inputs["output_router_logits"]
             )
+        if self.config.architectures[0] == "JambaForCausalLM":
+            example_dict["num_logits_to_keep"] = torch.tensor(0)
 
         with torch.inference_mode(), torch.no_grad(), torch.cpu.amp.autocast(
             enabled=True if args.quant_with_amp or self._dtype == "bfloat16" else False,

examples/cpu/llm/inference/single_instance/run_accuracy.py

@@ -431,6 +431,8 @@ def _model_call(
             example_dict["output_router_logits"] = torch.tensor(
                 model_inputs["output_router_logits"]
             )
+        if self.config.architectures[0] == "JambaForCausalLM":
+            example_dict["num_logits_to_keep"] = torch.tensor(0)
 
         with torch.inference_mode(), torch.no_grad(), torch.cpu.amp.autocast(
             enabled=True if args.quant_with_amp or self._dtype == "bfloat16" else False,