    def __init__(self, dim=0) -> None:
        """initialize the op"""
        super().__init__()
        self._attrs["op"] = "argmax"
        self._attrs["has_profiler"] = False
        self._attrs["dim"] = dim
        self._attrs["has_profiler"] = True
        self._attrs["workspace"] = 0
        self.exec_key_template = EXEC_KEY_TEMPLATE
initialize the op
python/aitemplate/compiler/ops/tensor/argmax.py