RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:0!

[haoran1062]

hello, I'm trying to modify some knowledge to Qwen-14b, I have 8 x A100 to use, in yaml file

• when model_parallel : False, OOM occured
• when I set model_parallel : True use multi-GPU to run ROME/MEMIT RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:0! error occured
• qwen-14b.yaml

alg_name: "ROME"
model_name: "/data/weights/trained/sft/qwen-14b_hf/"
stats_dir: "./data/stats"
device: 0
layers: [10]
fact_token: "subject_last"
v_num_grad_steps: 20
v_lr: 5e-1
v_loss_layer: 39
v_weight_decay: 0.5
clamp_norm_factor: 4
kl_factor: 0.0625
mom2_adjustment: false
context_template_length_params: [[5, 10], [10, 10]]
rewrite_module_tmp: "transformer.h.{}.mlp.c_proj"
layer_module_tmp: "transformer.h.{}"
mlp_module_tmp: "transformer.h.{}.mlp"
attn_module_tmp: "transformer.h.{}.attn"
ln_f_module: "transformer.ln_f"
lm_head_module: "lm_head"
mom2_dataset: "wikipedia"
mom2_n_samples: 100000
mom2_dtype: "float32"
model_parallel: true

• output

Rewrite layer is 10                                                                                                                                                                                                                                                                                                                                                        
Tying optimization objective to 39                                                                                                                                                                                                                                                                                                                                         
Recording initial value of v*                                                                                                                                                                                                                                                                                                                                              
Traceback (most recent call last):                                                                                                                                                                                                                                                                                                                                         
  File "/data/projects/EasyEdit/play_edit.py", line 88, in <module>                                                                                                                                                                                                                                                                                                        
    test_ROME_Qwen(cfg_path, model_path)                                                                                                                                                                                                                                                                                                                                   
  File "/data/projects/EasyEdit/play_edit.py", line 32, in test_ROME_Qwen                                                                                                                                                                                                                                                                                                  
    metrics, edited_model, _ = editor.edit(                                                                                                                                                                                                                                                                                                                                
  File "/data/projects/EasyEdit/easyeditor/editors/editor.py", line 247, in edit                                                                                                                                                                                                                                                                                           
    edited_model, weights_copy = self.apply_algo(                                                                                                                                                                                                                                                                                                                          
  File "/data/projects/EasyEdit/easyeditor/models/rome/rome_main.py", line 41, in apply_rome_to_model                                                                                                                                                                                                                                                                      
    deltas = execute_rome(model, tok, request, hparams)                                                                                                                                                                                                                                                                                                                    
  File "/data/projects/EasyEdit/easyeditor/models/rome/rome_main.py", line 113, in execute_rome                                                                                                                                                                                                                                                                            
    right_vector: torch.Tensor = compute_v(                                                                                                                                                                                                                                                                                                                                
  File "/data/projects/EasyEdit/easyeditor/models/rome/compute_v.py", line 121, in compute_v                                                                                                                                                                                                                                                                               
    logits = model(**input_tok).logits                                                                                                                                                                                                                                                                                                                                     
  File "/opt/conda/envs/EasyEdit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                                                                                                                                                                                         
    return forward_call(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                   
  File "/opt/conda/envs/EasyEdit/lib/python3.9/site-packages/accelerate/hooks.py", line 164, in new_forward                                                                                                                                                                                                                                                                
    output = module._old_forward(*args, **kwargs)                                                                                                                                                                                                                                                                                                                          
  File "/root/.cache/huggingface/modules/transformers_modules/modeling_qwen.py", line 1104, in forward                                                                                                                                                                                                                                                                     
    transformer_outputs = self.transformer(                                                                                                                                                                                                                                                                                                                                
  File "/opt/conda/envs/EasyEdit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                                                                                                                                                                                         
    return forward_call(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                   
  File "/root/.cache/huggingface/modules/transformers_modules/modeling_qwen.py", line 934, in forward                                                                                                                                                                                                                                                                      
    outputs = block(                                                                                                                                                                                                                                                                                                                                                       
  File "/opt/conda/envs/EasyEdit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                                                                                                                                                                                                         
    return forward_call(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                   
  File "/opt/conda/envs/EasyEdit/lib/python3.9/site-packages/accelerate/hooks.py", line 164, in new_forward                                                                                                                                                                                                                                                                
    output = module._old_forward(*args, **kwargs)                                                                                                                                                                                                                                                                                                                          
  File "/root/.cache/huggingface/modules/transformers_modules/modeling_qwen.py", line 655, in forward                                                                                                                                                                                                                                                                      
    mlp_output = self.mlp(layernorm_output)                                                                                                                                                                                                                                                                                                                                
  File "/opt/conda/envs/EasyEdit/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1547, in _call_impl                                                                                                                                                                                                                                                         
    hook_result = hook(self, args, result)                                                                                                                                                                                                                                                                                                                                 
  File "/data/projects/EasyEdit/easyeditor/util/nethook.py", line 80, in retain_hook                                                                                                                                                                                                                                                                                       
    output = invoke_with_optional_args(                                                                                                                                                                                                                                                                                                                                    
  File "/data/projects/EasyEdit/easyeditor/util/nethook.py", line 451, in invoke_with_optional_args                                                                                                                                                                                                                                                                        
    return fn(*pass_args, **pass_kw)                                                                                                                                                                                                                                                                                                                                       
  File "/data/projects/EasyEdit/easyeditor/models/rome/compute_v.py", line 98, in edit_output_fn                                                                                                                                                                                                                                                                           
    cur_out[i, idx, :] += delta                                                                                                                                                                                                                                                                                                                                            
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:0! 

[pengzju]

https://github.com/zjunlp/EasyEdit/blob/main/easyeditor/models/rome/compute_v.py#L81
You can set this delta to the device of layer 10

[zxlzr]

Hi, have you solved your issue?

[peixin-lin]

https://github.com/zjunlp/EasyEdit/blob/main/easyeditor/models/rome/compute_v.py#L81 You can set this delta to the device of layer 10

Is there a pythonic way to do this or we can only determine which device layer 10 will be on and set it in advance?

[littlefive5]

Provisionally, try

cur_out[i, idx, :] += delta.to(cur_out.device)

[JX-Zhang98]

https://github.com/zjunlp/EasyEdit/blob/main/easyeditor/models/rome/compute_v.py#L81 You can set this delta to the device of layer 10

I meet the same problem. I attempt to edit chatglm3-6b on multi-gpu, and my yaml is like:

alg_name: "ROME"
model_name: /home/gwz/glm3/chatglm3-6b
device: 1
mom2_dtype: "float16"
model_parallel: true

other options is the same with chatglm2-6b.yaml.

In my scenario, the rewrite layer is 5, and the edit.py prints:

Selected u projection object 中国
Left vector shape: torch.Size([13696])
Computing right vector (v)
Lookup index found: 2 | Sentence: 中国的首都是哪里？[gMASK]sop   | Token: 
Rewrite layer is 5
Tying optimization objective to 27
Recording initial value of v*

After trying, I determined that layer 5 is on device gpu1，because the error is different only when device is set to 'cuda:1'.
So I specified device in yaml as 1, but it meets the same error at a new location instead of the line of cur_out[i, idx, :] += delta :

rewrite layer is 5
Tying optimization objective to 27
Recording initial value of v*
Traceback (most recent call last):
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/main.py", line 107, in <module>
    metrics, edited_model, tokenizer = test_ROME_ChatGLM3()
                                       ^^^^^^^^^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/main.py", line 82, in test_ROME_ChatGLM3
    metrics, edited_model, _ = editor.edit(
                               ^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/easyeditor/editors/editor.py", line 252, in edit
    edited_model, weights_copy = self.apply_algo(
                                 ^^^^^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/easyeditor/models/rome/rome_main.py", line 41, in apply_rome_to_model
    deltas = execute_rome(model, tok, request, hparams)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/easyeditor/models/rome/rome_main.py", line 113, in execute_rome
    right_vector: torch.Tensor = compute_v(
                                 ^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/easyeditor/models/rome/compute_v.py", line 156, in compute_v
    loss = nll_loss + kl_loss + weight_decay
           ~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

I wonder what I need to do to run the scripts on multi-gpu, thank you~

[littlefive5]

If you set model_parallel as true, the device would not work here.
Provisionally, try

cur_out[i, idx, :] += delta.to(cur_out.device)

or

loss = nll_loss + kl_loss.to(nll_loss .device) + weight_decay.to(nll_loss .device)

[JX-Zhang98]

@littlefive5 Thanks for help.
It seems works when two gpus are given, though it still fails on more than two gpu...
Anyway, I got a milestone to update weights of the model. But a new challenge is comming:

Selected u projection object 中国
Left vector shape: torch.Size([13696])
Computing right vector (v)
Lookup index found: 2 | Sentence: 中国的首都是哪里？[gMASK]sop   | Token: 
Rewrite layer is 5
Tying optimization objective to 27
Recording initial value of v*
loss 19.545 = 19.545 + 0.0 + 0.0 avg prob of [ 纽约] 1.30896919969814e-08
loss 18.421 = 18.162 + 0.061 + 0.198 avg prob of [ 纽约] 3.3666495369288896e-08
loss 16.044 = 15.811 + 0.01 + 0.222 avg prob of [ 纽约] 5.215836722527456e-07
loss 13.512 = 13.279 + 0.011 + 0.222 avg prob of [ 纽约] 7.951258339744527e-06
loss 11.479 = 11.246 + 0.011 + 0.222 avg prob of [ 纽约] 2.8541777282953262e-05
loss 10.19 = 9.957 + 0.011 + 0.222 avg prob of [ 纽约] 8.683399937581271e-05
loss 8.939 = 8.706 + 0.011 + 0.222 avg prob of [ 纽约] 0.0003116014995612204
loss 7.684 = 7.451 + 0.011 + 0.222 avg prob of [ 纽约] 0.0012688582064583898
loss 6.708 = 6.475 + 0.011 + 0.222 avg prob of [ 纽约] 0.0040318709798157215
loss 6.432 = 6.199 + 0.011 + 0.222 avg prob of [ 纽约] 0.004579639527946711
loss 5.269 = 5.036 + 0.011 + 0.222 avg prob of [ 纽约] 0.016872501000761986
loss 4.886 = 4.653 + 0.011 + 0.222 avg prob of [ 纽约] 0.021900227293372154
loss 4.959 = 4.727 + 0.01 + 0.222 avg prob of [ 纽约] 0.017623111605644226
loss 5.239 = 5.007 + 0.009 + 0.222 avg prob of [ 纽约] 0.012865194119513035
loss 4.694 = 4.463 + 0.009 + 0.222 avg prob of [ 纽约] 0.02247370406985283
loss 4.415 = 4.184 + 0.009 + 0.222 avg prob of [ 纽约] 0.025304453447461128
loss 4.049 = 3.819 + 0.008 + 0.222 avg prob of [ 纽约] 0.03387002646923065
loss 3.861 = 3.632 + 0.007 + 0.222 avg prob of [ 纽约] 0.037830062210559845
loss 3.597 = 3.369 + 0.005 + 0.222 avg prob of [ 纽约] 0.04446116462349892
loss 3.413 = 3.187 + 0.003 + 0.222 avg prob of [ 纽约] 0.05012023076415062
Delta norm: 39.899391174316406
Change in target norm: 8.9921875 to 37.08768844604492 => 28.095500946044922
Division Factor: 0.83447265625
Right vector norm: 47.813899993896484
Right vector shape: torch.Size([4096])
Traceback (most recent call last):
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/main.py", line 107, in <module>
    metrics, edited_model, tokenizer = test_ROME_ChatGLM3()
                                       ^^^^^^^^^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/main.py", line 82, in test_ROME_ChatGLM3
    metrics, edited_model, _ = editor.edit(
                               ^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/easyeditor/editors/editor.py", line 252, in edit
    edited_model, weights_copy = self.apply_algo(
                                 ^^^^^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/easyeditor/models/rome/rome_main.py", line 41, in apply_rome_to_model
    deltas = execute_rome(model, tok, request, hparams)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zjx/workspace/AIrelated/AI-Evaluator/attack_verify/model_edit/easyeditor/models/rome/rome_main.py", line 127, in execute_rome
    upd_matrix = left_vector.unsqueeze(1) @ right_vector.unsqueeze(0)
                 ~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: c10::Half != float

Do you have any idea? Thank you~

[JX-Zhang98]

I debug the script and on line of upd_matrix = left_vector.unsqueeze(1) @ right_vector.unsqueeze(0), I get :

> left_vector.unsqueeze(1).shape
torch.Size([13696, 1])
> right_vector.unsqueeze(0).shape 
torch.Size([1, 4096])

It seems the col of left_vec.unsqueeze(1) is equals to the raw of right_vec.unsqueeze(0).

---

This problem has been solved. It occurred due to me adjusting the dtype when loading model, and can be solved by just aligning the two vectors
left_vector = left_vector.to(right_vector.dtype)

[littlefive5]

@JX-Zhang98
Glad to see your success running.
For multi-gpus, you can also move the tensors to the same device where the error is evoked as a temporary solution.

[zxlzr]

Hi, do you have any other issues?

[peixin-lin]

If you set model_parallel as true, the device would not work here. Provisionally, try

cur_out[i, idx, :] += delta.to(cur_out.device)

or

loss = nll_loss + kl_loss.to(nll_loss .device) + weight_decay.to(nll_loss .device)

When I use MEMIT for multi-GPUs, I had to modify multiple lines in easyeditor/models/memit/compute_z.py (In details, they are: line 106, line 108, line 151, line 166, line 187, line 192. ) to avoid this error. The base LM to be edited is Qwen-14B-Chat.
However, when I run this code I got OOM error although I used 8 * 80G H800 to run it. Did I make some mistake in my code? Please advise!

All I have done is to make to move tensors on the same device as you done in my opinion.
The complete code of compute_z.py after my edit is as follow:

from typing import Dict, List, Tuple

import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from ..rome import repr_tools
from ...util import nethook

from .memit_hparams import MEMITHyperParams


def compute_z(
    model: AutoModelForCausalLM,
    tok: AutoTokenizer,
    request: Dict,
    hparams: MEMITHyperParams,
    layer: int,
    context_templates: List[str],
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Computes the value (right) vector for the rank-1 update.
    Runs a simple optimization procedure.
    """

    # Get model parameters
    lm_w, ln_f = (
        nethook.get_parameter(model, f"{hparams.lm_head_module}.weight").T,
        nethook.get_module(model, hparams.ln_f_module),
    )
    try:
        lm_b = nethook.get_parameter(model, f"{hparams.lm_head_module}.bias")
    except LookupError as _:
        lm_b = next(model.parameters()).new_zeros(model.config.vocab_size)

    print("Computing right vector (v)")

    # Tokenize target into list of int token IDs
    target_ids = tok(request["target_new"], return_tensors="pt").to(f"cuda:{hparams.device}")[
        "input_ids"
    ][0]

    if target_ids[0] == tok.bos_token_id or target_ids[0] == tok.unk_token_id:
        target_ids = target_ids[1:]
    # Compile list of rewriting and KL x/y pairs
    rewriting_prompts, kl_prompts = [
        context.format(request["prompt"]) + tok.decode(target_ids[:-1])
        for context_types in context_templates
        for context in context_types
    ], ["{} is a"]
    all_prompts = rewriting_prompts + kl_prompts

    input_tok = tok(
        [prompt.format(request["subject"]) for prompt in all_prompts],
        return_tensors="pt",
        padding=True,
    ).to(f"cuda:{hparams.device}")

    # Compute rewriting targets
    rewriting_targets = torch.tensor(-100, device=f"cuda:{hparams.device}").repeat(
        len(rewriting_prompts), *input_tok["input_ids"].shape[1:]
    )
    for i in range(len(rewriting_prompts)):
        ex_len = input_tok["attention_mask"][i].sum()
        rewriting_targets[i, ex_len - len(target_ids) : ex_len] = target_ids

    # Compute indices of the tokens where the fact is looked up
    lookup_idxs = [
        find_fact_lookup_idx(
            prompt, request["subject"], tok, hparams.fact_token, verbose=(i == 0)
        )
        for i, prompt in enumerate(all_prompts)
    ]

    # Finalize rewrite and loss layers
    loss_layer = max(hparams.v_loss_layer, layer)
    print(f"Rewrite layer is {layer}")
    print(f"Tying optimization objective to {loss_layer}")

    # Set up an optimization over a latent vector that, when output at the
    # rewrite layer, i.e. hypothesized fact lookup location, will induce the
    # target token to be predicted at the final layer.
    if hasattr(model.config, 'n_embd'):
        delta = torch.zeros((model.config.n_embd,), requires_grad=True, device=f"cuda:{hparams.device}")
    elif hasattr(model.config, 'hidden_size'):
        delta = torch.zeros((model.config.hidden_size,), requires_grad=True, device=f"cuda:{hparams.device}")
    else:
        raise NotImplementedError
    target_init, kl_distr_init = None, None

    # Inserts new "delta" variable at the appropriate part of the computation
    def edit_output_fn(cur_out, cur_layer):
        nonlocal target_init

        if cur_layer == hparams.layer_module_tmp.format(layer):
            # Store initial value of the vector of interest
            if target_init is None:
                print("Recording initial value of v*")
                # Initial value is recorded for the clean sentence
                target_init = cur_out[0][0, lookup_idxs[0]].detach().clone()

            # Add intervened delta
            for i, idx in enumerate(lookup_idxs):

                if len(lookup_idxs)!=len(cur_out[0]):
                    cur_out[0][idx, i, :] += delta.to(cur_out[0].device)  # modified
                else:
                    cur_out[0][i, idx, :] += delta.to(cur_out[0].device)  # modified

        return cur_out

    # Optimizer
    opt = torch.optim.Adam([delta], lr=hparams.v_lr)
    nethook.set_requires_grad(False, model)

    # Execute optimization
    for it in range(hparams.v_num_grad_steps):
        opt.zero_grad()

        # Forward propagation
        with nethook.TraceDict(
            module=model,
            layers=[
                hparams.layer_module_tmp.format(loss_layer),
                hparams.layer_module_tmp.format(layer),
            ],
            retain_input=False,
            retain_output=True,
            edit_output=edit_output_fn,
        ) as tr:
            logits = model(**input_tok).logits
            # Compute distribution for KL divergence
            kl_logits = torch.stack(
                [
                    logits[i - len(kl_prompts), idx, :]
                    for i, idx in enumerate(lookup_idxs[-len(kl_prompts) :])
                ],
                dim=0,
            )
            kl_log_probs = torch.nn.functional.log_softmax(kl_logits, dim=1)
            if kl_distr_init is None:
                kl_distr_init = kl_log_probs.detach().clone()

        # Compute loss on rewriting targets

        output=tr[hparams.layer_module_tmp.format(loss_layer)].output[0]
        if output.shape[1]!=rewriting_targets.shape[1]:
            output=torch.transpose(output, 0, 1)
        full_repr = output[:len(rewriting_prompts)]

        log_probs = torch.log_softmax(ln_f(full_repr).to(full_repr.device) @ lm_w.to(full_repr.device) + lm_b.to(full_repr.device), dim=2)  # modified
        loss = torch.gather(
            log_probs,
            2,
            torch.where(rewriting_targets != -100, rewriting_targets, 0).unsqueeze(2).to(log_probs.device),
        ).squeeze(2)
        mask = (rewriting_targets != -100).float()

        # Aggregate total losses
        nll_loss_each = -(loss * mask.to(loss.device)).sum(1) / target_ids.size(0)
        nll_loss = nll_loss_each.mean()
        kl_loss = hparams.kl_factor * torch.nn.functional.kl_div(
            kl_distr_init, kl_log_probs, log_target=True, reduction="batchmean"
        )
        weight_decay = hparams.v_weight_decay * (
            torch.norm(delta) / torch.norm(target_init).to(delta.device) ** 2  
        )  # modified
        # weight_decay = hparams.v_weight_decay * torch.norm(delta) ** 2
        loss = nll_loss + kl_loss.to(nll_loss.device) + weight_decay.to(nll_loss.device)
        print(
            f"loss {np.round(loss.item(), 3)} = {np.round(nll_loss.item(), 3)} + {np.round(kl_loss.item(), 3)} + {np.round(weight_decay.item(), 3)} "
            f"avg prob of [{request['target_new']}] "
            f"{torch.exp(-nll_loss_each).mean().item()}"
        )
        if loss < 5e-2:
            break

        if it == hparams.v_num_grad_steps - 1:
            break

        # Backpropagate
        loss.backward()
        opt.step()

        # Project within L2 ball
        max_norm = hparams.clamp_norm_factor * target_init.norm()
        max_norm = max_norm.to(delta.device)  # modified
        if delta.norm() > max_norm:
            with torch.no_grad():
                delta[...] = delta * max_norm / delta.norm()

    target = target_init + delta.to(target_init.device)  # modified
    print(
        f"Init norm {target_init.norm()} | Delta norm {delta.norm()} | Target norm {target.norm()}"
    )

    return target


def get_module_input_output_at_words(
    model: AutoModelForCausalLM,
    tok: AutoTokenizer,
    layer: int,
    context_templates: List[str],
    words: List[str],
    module_template: str,
    fact_token_strategy: str,
    track=None,
) -> Tuple[torch.Tensor]:
    """
    Retrieves detached representations for a word at the input and
    output of a particular layer module.
    """

    word_repr_args = dict(
        model=model,
        tok=tok,
        layer=layer,
        module_template=module_template,
    )
    if "subject_" in fact_token_strategy and fact_token_strategy.index("subject_") == 0:
        context_info = dict(
            context_templates=context_templates,
            words=words,
        )
        subtoken = fact_token_strategy[len("subject_") :]
        if track == 'out' or track == 'in':
            return repr_tools.get_reprs_at_word_tokens(
                track=track, subtoken=subtoken, **context_info, **word_repr_args
            )
        l_input, l_output = repr_tools.get_reprs_at_word_tokens(
            track="both", subtoken=subtoken, **context_info, **word_repr_args
        )
    elif fact_token_strategy == "last":
        raise Exception("This is definitely bugged, fix it.")
        context_info = dict(
            contexts=[
                tmp[i].format(words[i]) for i, tmp in enumerate(context_templates)
            ],
            idxs=[000000],
        )
        if track == 'out' or track == 'in':
            return repr_tools.get_reprs_at_word_tokens(
                track=track, subtoken=subtoken, **context_info, **word_repr_args
            )
        l_input, l_output = repr_tools.get_reprs_at_idxs(
            track="both", **context_info, **word_repr_args
        )
    else:
        raise ValueError(f"fact_token={fact_token_strategy} not recognized")

    return l_input.detach(), l_output.detach()


def find_fact_lookup_idx(
    prompt: str,
    subject: str,
    tok: AutoTokenizer,
    fact_token_strategy: str,
    verbose=True,
) -> int:
    """
    Computes hypothesized fact lookup index given a sentence and subject.
    """

    ret = None
    if fact_token_strategy == "last":
        ret = -1
    elif (
        "subject_" in fact_token_strategy and fact_token_strategy.index("subject_") == 0
    ):
        ret = repr_tools.get_words_idxs_in_templates(
            tok=tok,
            context_templates=[prompt],
            words=[subject],
            subtoken=fact_token_strategy[len("subject_") :],
        )[0][0]
    else:
        raise ValueError(f"fact_token={fact_token_strategy} not recognized")

    sentence = prompt.format(subject)
    if verbose:
        print(
            f"Lookup index found: {ret} | Sentence: {sentence} | Token:",
            tok.decode(tok(sentence)["input_ids"][ret]),
        )

    return ret