def METHOD_NAME(self, model: ModelProto) -> ModelProto:
    """
    Adds causal mask as an input to the model.
    Causal mask is an int64 tensor of shape
    [batch_size, 1, input_ids_length, sequence_length]
    where the value is 0 if the position is masked and 1
    otherwise.
    :param model: model to update
    :return: updated model
    """
    input_ids = self._get_input_proto(model, "input_ids")
    attention_mask = self._get_input_proto(model, "attention_mask")
    batch_size = input_ids.type.tensor_type.shape.dim[0].dim_param
    input_ids_length = input_ids.type.tensor_type.shape.dim[1].dim_value
    sequence_length = attention_mask.type.tensor_type.shape.dim[1].dim_param
    causal_mask_input = helper.make_tensor_value_info(
        name=self.CAUSAL_MASK_NAME,
        elem_type=TensorProto.INT64,
        shape=[batch_size, 1, input_ids_length, sequence_length],
    )
    model.graph.input.append(causal_mask_input)
    _LOGGER.info(f"Inserted {self.CAUSAL_MASK_NAME} input to the ONNX model")
    return model