import torch
import operator


def change_attention_mask_value(traced, initial_value=-10000, final_value=-20):
    """Changes the attention mask initial value (default is -10000) to some smaller value."""
    graph = traced.graph
    for node in graph.nodes:
        if node.target in [torch.mul, operator.mul] and initial_value in node.args:
            new_args = []
            for arg in node.args:
                if arg == initial_value:
                    new_args.append(final_value)
                else:
                    new_args.append(arg)
            node.args = tuple(new_args)
            break

    graph.lint()
    traced.recompile()
    return traced