ONNX Conversion error

[kishoreb4]

I was able follow the script and finetune the model of a new dataset for forecasting.However, when i convert to onnx i'm getting the below error. Any help would be appreciated. Thanks in advance.

im getting the below error when in the code torch.onnx.export(model, (timeseries, input_mask), onnx_file_path)

---

RuntimeError Traceback (most recent call last)
Cell In[59], line 2
1 # Convert the model to ONNX
----> 2 torch.onnx.export(model, (timeseries, input_mask), onnx_file_path)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:516, in export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, custom_opsets, export_modules_as_functions, autograd_inlining)
189 @_beartype.beartype
190 def export(
191 model: Union[torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction],
(...)
208 autograd_inlining: Optional[bool] = True,
209 ) -> None:
210 r"""Exports a model into ONNX format.
211
212 If model is not a :class:torch.jit.ScriptModule nor a
(...)
513 All errors are subclasses of :class:errors.OnnxExporterError.
514 """
--> 516 _export(
517 model,
518 args,
519 f,
520 export_params,
521 verbose,
522 training,
523 input_names,
524 output_names,
525 operator_export_type=operator_export_type,
526 opset_version=opset_version,
527 do_constant_folding=do_constant_folding,
528 dynamic_axes=dynamic_axes,
529 keep_initializers_as_inputs=keep_initializers_as_inputs,
530 custom_opsets=custom_opsets,
531 export_modules_as_functions=export_modules_as_functions,
532 autograd_inlining=autograd_inlining,
533 )

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:1612, in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, onnx_shape_inference, export_modules_as_functions, autograd_inlining)
1609 dynamic_axes = {}
1610 _validate_dynamic_axes(dynamic_axes, model, input_names, output_names)
-> 1612 graph, params_dict, torch_out = _model_to_graph(
1613 model,
1614 args,
1615 verbose,
1616 input_names,
1617 output_names,
1618 operator_export_type,
1619 val_do_constant_folding,
1620 fixed_batch_size=fixed_batch_size,
1621 training=training,
1622 dynamic_axes=dynamic_axes,
1623 )
1625 # TODO: Don't allocate a in-memory string for the protobuf
1626 defer_weight_export = (
1627 export_type is not _exporter_states.ExportTypes.PROTOBUF_FILE
1628 )

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:1134, in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)
1131 args = (args,)
1133 model = _pre_trace_quant_model(model, args)
-> 1134 graph, params, torch_out, module = _create_jit_graph(model, args)
1135 params_dict = _get_named_param_dict(graph, params)
1137 try:

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:1010, in _create_jit_graph(model, args)
1005 graph = _C._propagate_and_assign_input_shapes(
1006 graph, flattened_args, param_count_list, False, False
1007 )
1008 return graph, params, torch_out, None
-> 1010 graph, torch_out = _trace_and_get_graph_from_model(model, args)
1011 _C._jit_pass_onnx_lint(graph)
1012 state_dict = torch.jit._unique_state_dict(model)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/onnx/utils.py:914, in _trace_and_get_graph_from_model(model, args)
912 prev_autocast_cache_enabled = torch.is_autocast_cache_enabled()
913 torch.set_autocast_cache_enabled(False)
--> 914 trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
915 model,
916 args,
917 strict=False,
918 _force_outplace=False,
919 _return_inputs_states=True,
920 )
921 torch.set_autocast_cache_enabled(prev_autocast_cache_enabled)
923 warn_on_static_input_change(inputs_states)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/jit/_trace.py:1310, in _get_trace_graph(f, args, kwargs, strict, _force_outplace, return_inputs, _return_inputs_states)
1308 if not isinstance(args, tuple):
1309 args = (args,)
-> 1310 outs = ONNXTracedModule(
1311 f, strict, _force_outplace, return_inputs, _return_inputs_states
1312 )(*args, **kwargs)
1313 return outs

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
1530 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1531 else:
-> 1532 return self._call_impl(*args, **kwargs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
1536 # If we don't have any hooks, we want to skip the rest of the logic in
1537 # this function, and just call forward.
1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1539 or _global_backward_pre_hooks or _global_backward_hooks
1540 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541 return forward_call(*args, **kwargs)
1543 try:
1544 result = None

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/jit/_trace.py:138, in ONNXTracedModule.forward(self, *args)
135 else:
136 return tuple(out_vars)
--> 138 graph, out = torch._C._create_graph_by_tracing(
139 wrapper,
140 in_vars + module_state,
141 _create_interpreter_name_lookup_fn(),
142 self.strict,
143 self._force_outplace,
144 )
146 if self._return_inputs:
147 return graph, outs[0], ret_inputs[0]

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/torch/jit/_trace.py:129, in ONNXTracedModule.forward..wrapper(*args)
127 if self._return_inputs_states:
128 inputs_states.append(_unflatten(in_args, in_desc))
--> 129 outs.append(self.inner(*trace_inputs))
130 if self._return_inputs_states:
131 inputs_states[0] = (inputs_states[0], trace_inputs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/momentfm/models/moment.py:222, in MOMENT.call(self, *args, **kwargs)
221 def call(self, *args, **kwargs) -> TimeseriesOutputs:
--> 222 return self.forward(*args, **kwargs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/momentfm/models/moment.py:549, in MOMENT.forward(self, x_enc, mask, input_mask, **kwargs)
547 return self.embed(x_enc=x_enc, input_mask=input_mask, **kwargs)
548 elif self.task_name == TASKS.FORECASTING:
--> 549 return self.forecast(x_enc=x_enc, input_mask=input_mask, **kwargs)
550 elif self.task_name == TASKS.CLASSIFICATION:
551 return self.classify(x_enc=x_enc, input_mask=input_mask, **kwargs)

File /libraries/gpu_py_3_11/lib/python3.11/site-packages/momentfm/models/moment.py:423, in MOMENT.forecast(self, x_enc, input_mask, **kwargs)
418 enc_in = enc_in.reshape(
419 (batch_size * n_channels, n_patches, self.config.d_model)
420 )
422 patch_view_mask = Masking.convert_seq_to_patch_view(input_mask, self.patch_len)
--> 423 attention_mask = patch_view_mask.repeat_interleave(n_channels, dim=0)
424 outputs = self.encoder(inputs_embeds=enc_in, attention_mask=attention_mask)
425 enc_out = outputs.last_hidden_state

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)