Export to ONNX supported?

[dnth]

I'm trying to export the model into .onnx.

Here's my code:

import torch
model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14').to('cuda:0')
model.eval()

# Generate some input data
input_data = torch.randn(1, 3, 224, 224).to('cuda:0')

# Pass the input data through the model
output = model(input_data)

torch.onnx.export(model, input_data, 'model.onnx')

I got an error

============= Diagnostic Run torch.onnx.export version 2.0.0+cu117 =============
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-17-2f8453f4374c>](https://localhost:8080/#) in <cell line: 1>()
----> 1 torch.onnx.export(model, input_data, 'model.onnx')

13 frames
[~/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/models/vision_transformer.py](https://localhost:8080/#) in prepare_tokens_with_masks(self, x, masks)
    193         x = self.patch_embed(x)
    194         if masks is not None:
--> 195             x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
    196 
    197         x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Did I miss anything?