    def _infer_shape(self, x: List[int], w: List[int]):
        """Infer the output shape"""
        eval_func = self.shape_eval_template.render(
            indent="",
            dtype="",
            div="//",
            nmsMaxOut=self._attrs["nmsMaxOut"],
            x_dim0=x[0],
            x_dim1=x[1],
            x_dim2=x[2],
        )
        output = {}
        exec(eval_func, output)  # noqa: P204  # noqa: P204
        return [int(output["NO"]), int(output["CO"]), int(output["HO"])]
Infer the output shape
python/aitemplate/compiler/ops/vision_ops/nms/nms.py