fix

wenxindongwork · wenxindongwork · commit 3719a86fce57 · 2025-04-24T18:44:14.000Z
diff --git a/jetstream/engine/engine_api.py b/jetstream/engine/engine_api.py
@@ -155,7 +155,7 @@ def get_result_at_slots(self, slots: tuple[int]) -> SlotData:
         valid=self.data[slots, self.valid_idx[0] : self.valid_idx[1]],
         # Only get a 1D representation here
         lengths=self.data[slots, self.length_idx[0] : self.length_idx[1]][:, 0],
-        log_prob=self.log_prob[slots, :],
+        log_prob=self.log_prob[slots, :] if self.log_prob is not None else None,
     )
 
 

Original file line number	Diff line number	Diff line change
`@@ -155,7 +155,7 @@ def get_result_at_slots(self, slots: tuple[int]) -> SlotData:`
`155`	`155`	`valid=self.data[slots, self.valid_idx[0] : self.valid_idx[1]],`
`156`	`156`	`# Only get a 1D representation here`
`157`	`157`	`lengths=self.data[slots, self.length_idx[0] : self.length_idx[1]][:, 0],`
`158`		`- log_prob=self.log_prob[slots, :],`
	`158`	`+ log_prob=self.log_prob[slots, :] if self.log_prob is not None else None,`
`159`	`159`	`)`
`160`	`160`
`161`	`161`