Explicit device assignment ignored in multi-GPU (PyTriton) environment

I am running multiple SentenceTransformer models in a single Python process (via PyTriton). My goal is to orchestrate different models onto different physical GPUs (e.g., cuda:0, cuda:1).

Despite passing an explicit device argument (e.g., device='cuda:1') and attempting to use torch.cuda.set_device() context managers, I observe significant memory allocation on cuda:0 (the default device) that scales with the model size, suggesting the model weights or their initialization stages are being staged on the wrong GPU before being moved.

Environment:

sentence-transformers version: 5.5.1

torch version: 2.12.0

OS: Ubuntu 22.04

Expected Behavior: All model weights and internal buffers should be allocated exclusively on the target GPU provided in the device parameter.

Actual Behavior: The models consistently allocate memory on cuda:0. I have verified this using nvidia-smi. The memory footprint on cuda:0 matches the size of the model weights, confirming this is not just the standard "Global CUDA Context" overhead, but a genuine staging issue.

What I have already tried: To isolate the issue, I have attempted the following:

• Manual Module Injection: Initializing models.Transformer and models.Pooling separately and moving them .to(device) before passing them to the SentenceTransformer constructor.

• Hardware Lockdown: Using torch.cuda.set_device(gpu_id) context managers during the initialization phase.

• RAM-First Loading: Loading the model with device='cpu' first, then calling .to(device) before the first inference.

• model_kwargs: Injecting device_map via AutoModel.from_pretrained (when bypassing the SentenceTransformer constructor).

none of these prevent the initial weight staging/allocation on cuda:0.