def METHOD_NAME(self, module, inputs):
    """Pass initial activations to graph generation pass"""
    device = inputs[0].device if isinstance(inputs, tuple) else inputs.device
    for input, activation in zip(inputs, module.activations[device]):
        input.data = activation
    return inputs