Keep input_tokens on the model device in generate()

[danielxmed]

Description

HookedTransformer.generate() crashes on CUDA machines when the input is a token tensor on a different device than the model (e.g. CPU tokens passed to a model whose cfg.device is "cuda", the default when CUDA is available):

RuntimeError: Expected all tensors to be on the same device, but got tensors is on cuda:0, different from other tensors on cpu (when checking argument in method wrapper_CUDA_cat)

Root cause: input_tokens is bound to the original input tensor before input = input.to(device) moves the input to the model's device, so it keeps pointing at the original CPU tensor. Sampled tokens are produced on the model's device, so torch.cat((input_tokens, ...)) mixes devices — both in the freq-penalty sampling path and in the final output concatenation (the latter also hits do_sample=False). The ordering came in with #999.

Fix: after input = input.to(device), re-alias input_tokens = input (guarded by input_tokens is not None, which is exactly the str/tokens input case). Re-aliasing rather than calling .to(device) again avoids a second host-to-device copy of the same data.

This is easiest to hit with tokenizer-free models (algorithmic tasks), which is why 4 of the 6 existing tests in tests/unit/test_generate_no_tokenizer.py fail on any CUDA-visible machine at current dev while passing on CPU-only CI. With this fix they pass on both. generate_stream() and TransformerBridge.generate() are not affected (both already move their token tensor before use), so no Bridge mirroring is needed.

Tests: added test_generate_cpu_tokens_cuda_model, a CUDA-gated (skipif) regression test that explicitly passes CPU tokens to a CUDA model and covers both the sampling and greedy paths. Verified it fails without the fix and passes with it.

Ran locally on a CUDA machine (RTX 5070, torch 2.10.0+cu128, Python 3.12):

• pytest tests/unit/test_generate_no_tokenizer.py — 4 failed / 2 passed before, 7 passed after
• pytest tests/unit -m "not slow" — passes
• mypy transformer_lens/HookedTransformer.py, black/isort/pycln checks — clean

Fixes #1508

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility