diff --git a/selfdrive/modeld/parse_model_outputs.py b/selfdrive/modeld/parse_model_outputs.py index b89de27594..9e1c048735 100644 --- a/selfdrive/modeld/parse_model_outputs.py +++ b/selfdrive/modeld/parse_model_outputs.py @@ -94,15 +94,21 @@ class Parser: self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) self.parse_binary_crossentropy('meta', outs) self.parse_binary_crossentropy('lead_prob', outs) - self.parse_mdn('lead', outs, in_N=ModelConstants.LEAD_MHP_N, out_N=ModelConstants.LEAD_MHP_SELECTION, - out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH)) + if outs['lead'].shape[1] == 2 * ModelConstants.LEAD_MHP_SELECTION *ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH: + self.parse_mdn('lead', outs, in_N=0, out_N=0, + out_shape=(ModelConstants.LEAD_MHP_SELECTION, ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH)) + else: + self.parse_mdn('lead', outs, in_N=ModelConstants.LEAD_MHP_N, out_N=ModelConstants.LEAD_MHP_SELECTION, + out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH)) return outs def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - self.parse_mdn('plan', outs, in_N=ModelConstants.PLAN_MHP_N, out_N=ModelConstants.PLAN_MHP_SELECTION, - out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH)) - if 'lat_planner_solution' in outs: - self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N,ModelConstants.LAT_PLANNER_SOLUTION_WIDTH)) + if outs['plan'].shape[1] == 2 * ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH: + self.parse_mdn('plan', outs, in_N=0, out_N=0, + out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH)) + else: + self.parse_mdn('plan', outs, in_N=ModelConstants.PLAN_MHP_N, out_N=ModelConstants.PLAN_MHP_SELECTION, + out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH)) if 'desired_curvature' in outs: self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,)) self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,))