def METHOD_NAME(self, dataset):
    def point_clouds_only(point_clouds, target):
        return point_clouds
    def boxes_only(point_clouds, target):
        return target["3d_boxes"]
    model_outputs = self.model.predict(dataset.map(point_clouds_only))[
        "3d_boxes"
    ]
    def flatten_target(boxes):
        return tf.concat(
            [
                boxes["boxes"],
                tf.expand_dims(
                    tf.cast(boxes["classes"], tf.float32), axis=-1
                ),
                tf.expand_dims(
                    tf.cast(boxes["difficulty"], tf.float32), axis=-1
                ),
            ],
            axis=-1,
        )
    gt_boxes = tf.concat(
        [flatten_target(x) for x in iter(dataset.map(boxes_only))], axis=0
    )
    boxes_per_gt_frame = gt_boxes.shape[1]
    num_frames = gt_boxes.shape[0]
    gt_boxes = tf.reshape(gt_boxes, (num_frames * boxes_per_gt_frame, 9))
    # Remove padded boxes
    gt_real_boxes = tf.concat(
        [x["mask"] for x in iter(dataset.map(boxes_only))], axis=0
    )
    gt_real_boxes = tf.reshape(
        gt_real_boxes, (num_frames * boxes_per_gt_frame)
    )
    gt_boxes = tf.boolean_mask(gt_boxes, gt_real_boxes)
    frame_ids = tf.cast(tf.linspace(1, num_frames, num_frames), tf.int64)
    ground_truth = {
        "ground_truth_frame_id": tf.boolean_mask(
            tf.repeat(frame_ids, boxes_per_gt_frame), gt_real_boxes
        ),
        "ground_truth_bbox": gt_boxes[:, : CENTER_XYZ_DXDYDZ_PHI.PHI + 1],
        "ground_truth_type": tf.cast(
            gt_boxes[:, CENTER_XYZ_DXDYDZ_PHI.CLASS], tf.uint8
        ),
        "ground_truth_difficulty": tf.cast(
            gt_boxes[:, CENTER_XYZ_DXDYDZ_PHI.CLASS + 1], tf.uint8
        ),
    }
    boxes_per_pred_frame = model_outputs["boxes"].shape[1]
    total_predicted_boxes = boxes_per_pred_frame * num_frames
    predicted_boxes = tf.reshape(
        model_outputs["boxes"], (total_predicted_boxes, 7)
    )
    predicted_classes = tf.cast(
        tf.reshape(model_outputs["classes"], (total_predicted_boxes, 1)),
        tf.uint8,
    )
    prediction_scores = tf.reshape(
        model_outputs["confidence"], (total_predicted_boxes, 1)
    )
    # Remove boxes that come from padding
    pred_real_boxes = tf.squeeze(prediction_scores > 0)
    predicted_boxes = tf.boolean_mask(predicted_boxes, pred_real_boxes)
    predicted_classes = tf.boolean_mask(predicted_classes, pred_real_boxes)
    prediction_scores = tf.boolean_mask(prediction_scores, pred_real_boxes)
    predictions = {
        "prediction_frame_id": tf.boolean_mask(
            tf.repeat(frame_ids, boxes_per_pred_frame), pred_real_boxes
        ),
        "prediction_bbox": predicted_boxes,
        "prediction_type": tf.squeeze(predicted_classes),
        "prediction_score": tf.squeeze(prediction_scores),
        "prediction_overlap_nlz": tf.cast(
            tf.zeros(predicted_boxes.shape[0]), tf.bool
        ),
    }
    return ground_truth, predictions