Tensors on Different Devices During TorchScript Export

Hello, I am currently modifying the architecture from "[ONNX-SAM2-Segment-Anything](https://github.com/ibaiGorordo/ONNX-SAM2-Segment-Anything)" to export it to TorchScript. However, I have encountered a device mismatch issue during the export process. Below is my code:
```python
from typing import Optional, Tuple, Any
import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.init import trunc_normal_

from sam2.modeling.sam2_base import SAM2Base

class SAM2ImageEncoder(nn.Module):
def __init__(self, sam_model: SAM2Base) -> None:
super().__init__()
self.model = sam_model
self.image_encoder = sam_model.image_encoder
self.no_mem_embed = sam_model.no_mem_embed

def forward(self, x: torch.Tensor) -> tuple[Any, Any, Any]:
backbone_out = self.image_encoder(x)
backbone_out["backbone_fpn"][0] = self.model.sam_mask_decoder.conv_s0(
backbone_out["backbone_fpn"][0]
)
backbone_out["backbone_fpn"][1] = self.model.sam_mask_decoder.conv_s1(
backbone_out["backbone_fpn"][1]
)

feature_maps = backbone_out["backbone_fpn"][-self.model.num_feature_levels:]
vision_pos_embeds = backbone_out["vision_pos_enc"][-self.model.num_feature_levels:]

feat_sizes = [(x.shape[-2], x.shape[-1]) for x in vision_pos_embeds]

vision_feats = [x.flatten(2).permute(2, 0, 1) for x in feature_maps]
vision_pos_embeds = [x.flatten(2).permute(2, 0, 1) for x in vision_pos_embeds]

vision_feats[-1] = vision_feats[-1] + self.no_mem_embed

feats = [feat.permute(1, 2, 0).reshape(1, -1, *feat_size)
for feat, feat_size in zip(vision_feats[::-1], feat_sizes[::-1])][::-1]

return feats[0], feats[1], feats[2]

class SAM2ImageDecoder(nn.Module):
def __init__(self, sam_model: SAM2Base, multimask_output: bool) -> None:
super().__init__()
self.mask_decoder = sam_model.sam_mask_decoder
self.prompt_encoder = sam_model.sam_prompt_encoder
self.model = sam_model
self.multimask_output = multimask_output

@torch.no_grad()
def forward(self, image_embed: torch.Tensor, high_res_feats_0: torch.Tensor, high_res_feats_1: torch.Tensor, point_coords: torch.Tensor, point_labels: torch.Tensor, mask_input: torch.Tensor, has_mask_input: torch.Tensor, img_size: torch.Tensor):
sparse_embedding = self._embed_points(point_coords, point_labels)
dense_embedding = self._embed_masks(mask_input, has_mask_input)

high_res_feats = [high_res_feats_0, high_res_feats_1]

masks, iou_predictions, _, _ = self.mask_decoder.predict_masks(
image_embeddings=image_embed,
image_pe=self.prompt_encoder.get_dense_pe(),
sparse_prompt_embeddings=sparse_embedding,
dense_prompt_embeddings=dense_embedding,
repeat_image=False,
high_res_features=high_res_feats,
)

if self.multimask_output:
masks = masks[:, 1:, :, :]
iou_predictions = iou_predictions[:, 1:]
else:
masks, iou_predictions = self.mask_decoder._dynamic_multimask_via_stability(masks, iou_predictions)

masks = torch.clamp(masks, -32.0, 32.0)
print(masks.shape, iou_predictions.shape)

masks = F.interpolate(masks, (img_size[0], img_size[1]), mode="bilinear", align_corners=False)

return masks, iou_predictions

def _embed_points(self, point_coords: torch.Tensor, point_labels: torch.Tensor) -> torch.Tensor:
point_coords = point_coords + 0.5
padding_point = torch.zeros((point_coords.shape[0], 1, 2), device=point_coords.device)
padding_label = -torch.ones((point_labels.shape[0], 1), device=point_labels.device)
point_coords = torch.cat([point_coords, padding_point], dim=1)
point_labels = torch.cat([point_labels, padding_label], dim=1)

point_coords[:, :, 0] = point_coords[:, :, 0] / self.model.image_size
point_coords[:, :, 1] = point_coords[:, :, 1] / self.model.image_size

point_embedding = self.prompt_encoder.pe_layer._pe_encoding(point_coords)
point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding)

point_embedding = point_embedding * (point_labels != -1)
point_embedding = point_embedding + self.prompt_encoder.not_a_point_embed.weight * (
point_labels == -1
)

for i in range(self.prompt_encoder.num_point_embeddings):
point_embedding = point_embedding + self.prompt_encoder.point_embeddings[i].weight * (point_labels == i)

return point_embedding

def _embed_masks(self, input_mask: torch.Tensor, has_mask_input: torch.Tensor) -> torch.Tensor:
mask_embedding = has_mask_input * self.prompt_encoder.mask_downscaling(input_mask)
mask_embedding = mask_embedding + (
1 - has_mask_input
) * self.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)
return mask_embedding
```
```python
import traceback
from sam2.build_sam import build_sam2
def main():
try:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
model_type = 'sam2_hiera_large'
input_size = 1024
multimask_output = False

if model_type == "sam2_hiera_tiny":
model_cfg = "sam2_hiera_t.yaml"
elif model_type == "sam2_hiera_small":
model_cfg = "sam2_hiera_s.yaml"
elif model_type == "sam2_hiera_base_plus":
model_cfg = "sam2_hiera_b+.yaml"
else:
model_cfg = "sam2_hiera_l.yaml"

sam2_checkpoint = f"checkpoints/{model_type}.pt"
sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)

img = torch.randn(1, 3, input_size, input_size).to(device)

sam2_encoder = SAM2ImageEncoder(sam2_model).to(device)
high_res_feats_0, high_res_feats_1, image_embed = sam2_encoder(img)
#print(high_res_feats_0.shape)
#print(high_res_feats_1.shape)
#print(image_embed.shape)

sam2_decoder = SAM2ImageDecoder(sam2_model, multimask_output=multimask_output).to(device)

backbone_stride = sam2_model.backbone_stride
embed_size = (input_size // backbone_stride, input_size // backbone_stride)
mask_input_size = [4 * x for x in embed_size]

point_coords = torch.randint(low=0, high=input_size, size=(1, 5, 2), dtype=torch.float).to(device)
point_labels = torch.randint(low=0, high=1, size=(1, 5), dtype=torch.float).to(device)
mask_input = torch.randn(1, 1, *mask_input_size, dtype=torch.float).to(device) # 正確形狀的 mask_input
has_mask_input = torch.tensor([1], dtype=torch.float).to(device)
orig_im_size = torch.tensor([input_size, input_size], dtype=torch.int32).to(device)

traced_decoder = torch.jit.trace(sam2_decoder, (image_embed, high_res_feats_0, high_res_feats_1, point_coords, point_labels, mask_input, has_mask_input, orig_im_size))
traced_decoder.save("decoder.pt")

except Exception as e:
print(f"An error occurred: {e}")
traceback.print_exc()
if __name__ == "__main__":
main()
```
While trying to export the model using torch.jit.trace in a Colab environment with T4 GPU, I encountered the following error message:
```
cuda
An error occurred: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)
/content/segment-anything-2/sam2/modeling/sam/mask_decoder.py:203: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
assert image_embeddings.shape[0] == tokens.shape[0]
/content/segment-anything-2/sam2/modeling/sam/mask_decoder.py:207: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
image_pe.size(0) == 1
Traceback (most recent call last):
File "", line 43, in main
traced_decoder = torch.jit.trace(sam2_decoder, (image_embed, high_res_feats_0, high_res_feats_1, point_coords, point_labels, mask_input, has_mask_input, orig_im_size))
File "/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py", line 1000, in trace
traced_func = _trace_impl(
File "/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py", line 695, in _trace_impl
return trace_module(
File "/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py", line 1275, in trace_module
module._c._create_method_from_trace(
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1543, in _slow_forward
result = self.forward(*input, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "", line 58, in forward
masks, iou_predictions, _, _ = self.mask_decoder.predict_masks(
File "/content/segment-anything-2/sam2/modeling/sam/mask_decoder.py", line 209, in predict_masks
pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)
```
It appears that some tensors are on cpu while others are on cuda:0, even though I have ensured that all input tensors are moved to cuda. I suspect that some internal layers or tensors inside the mask_decoder are not being moved to cuda correctly. Despite moving all model layers to cuda, this issue persists.

I would greatly appreciate any suggestions or advice on how to resolve this device mismatch issue.

Best regards.