Fix: preserve device index in CrossLayerTranscoder lazy loading

[Xwwan]

What is the current behavior?

When initializing CrossLayerTranscoder with a specific device index (e.g., device="cuda:1") AND enabling lazy_decoder or lazy_encoder, the sub-modules are loaded onto the wrong device.
This happens because the code was accessing self.device.type (which returns just "cuda") instead of the full device object or string. This strips away the device index information, causing tensors/models to fall back to the default device (usually cuda:0).

What is the new behavior?

I updated the logic to use self.device (or the correct device string) during the lazy loading process. This ensures that if a user specifies cuda:1, the encoder/decoder is correctly initialized on cuda:1.

Steps to reproduce

import torch
from circuit_tracer.transcoder.cross_layer_transcoder import load_clt
from transformers import AutoModelForCausalLM, AutoTokenizer
from circuit_tracer.replacement_model import ReplacementModel
from circuit_tracer.attribution.attribute import attribute

model_path = "/PATH/TO/gemma-2-2b"
clt_path = "/PATH/TO/clt-gemma-2-2b-426k"
device = torch.device("cuda:1")
dtype = torch.bfloat16

clt = load_clt(
    clt_path,
    device=device,
    dtype=dtype,
    lazy_decoder=True,
    lazy_encoder=True,
)


hf_model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=dtype,
)

tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.padding_side = "left"
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id


model = ReplacementModel.from_pretrained_and_transcoders(
    model_name="google/gemma-2-2b",
    transcoders=clt,
    hf_model=hf_model,
    tokenizer=tokenizer,
    device=device,
    dtype=dtype,
    default_prepend_bos=True,
    default_padding_side="left",
)

prompt = "The capital of France is Paris."

graph = attribute(
    model=model,
    prompt=prompt,
    max_n_logits=5,
    desired_logit_prob=0.9,
    batch_size=64,
    max_feature_nodes=100,
    offload=None,
    verbose=False,
)

# Run Result: RuntimeError: Expected all tensors to be on the same device, but got mat2 is on cuda:0, different from other tensors on cuda:1 (when checking argument in method wrapper_CUDA_bmm)

[hannamw]

Thanks for opening this PR - this seems like a real issue! Would you mind making a couple of changes?:

• Instead of using the pattern

with safe_open(enc_file, framework="pt", device=self.device.type) as f:
                return f.get_tensor(f"W_enc_{layer_id}").to(dtype=self.dtype, device=self.device)

could you instead use this?

with safe_open(enc_file, framework="pt", device=str(self.device)) as f:
                return f.get_tensor(f"W_enc_{layer_id}").to(dtype)

This way, it'll just be loaded directly onto the correct device.

• Could you apply these changes to single_layer_transcoder.py as well? It suffers from the same issues.

Sorry to take so long to respond to this! Happy to do this myself as well if you don't have the time, but I thought I'd ask first