[BUG] embed function not working when sending cuda tensors

[HubertPalo]

Bug report checklist

• I provided code that demonstrates a minimal reproducible example.
• I confirmed bug exists on the latest mainline of Chronos via source install.

Describe the bug

I received an error when I used the embed function to extract embeddings from tensors in cuda.
The bug happens in all Chronos versions.
The cpu-based models are unaffected.
Tested on chronos-forecasting==2.3.1

Expected behavior

A similar behavior as when using a cpu tensor: the function passes the data through the model and returns the embeddings.

To reproduce

from chronos import Chronos2Pipeline
import torch
pipeline = Chronos2Pipeline.from_pretrained(
    "amazon/chronos-2",
    device_map="cuda",  # use "cpu" for CPU inference and "mps" for Apple Silicon
)
random_data = torch.rand(120, 6, 60).to("cuda")
pipeline.embed(random_data)[0][0].shape

However, this works:

from chronos import Chronos2Pipeline
import torch
pipeline = Chronos2Pipeline.from_pretrained(
    "amazon/chronos-2",
    device_map="cuda",  # use "cpu" for CPU inference and "mps" for Apple Silicon
)
random_data = torch.rand(120, 6, 60)
pipeline.embed(random_data)[0][0].shape

Environment description
Operating system: Ubuntu
Python version: 3.10.6
CUDA version: 13.2
PyTorch version: 2.9.1
HuggingFace transformers version:
HuggingFace accelerate version:
Pandas version: 2.3.3

[kaizokusk]

I was able to identify the root cause. When the model is on CUDA, embed() sets pin_memory=True even when the input tensor is already on CUDA. PyTorch only allows pinning dense CPU tensors, so this results in RuntimeError: cannot pin 'torch.cuda.FloatTensor'.

The fix is to enable pin_memory only when the input is on CPU and the model is on CUDA.