ONXX export issue Expected all tensors to be on the same device, but found at least two devices

[skagrawal10]

While trying to export Faster-VoxelPose model into onxx format, we are getting following issue -

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)

Details -

# Wrap the backbone and main model into a single module
class CombinedModel(torch.nn.Module):
    def __init__(self, backbone, model):
        super(CombinedModel, self).__init__()
        self.backbone = backbone
        self.model = model # Pretrained Faster voxel pose model

    def forward(self, views, meta, cameras, resize_transform):
        # Pass the backbone output and other inputs to the main model
        fused_poses, plane_poses, proposal_centers, input_heatmaps, _ = self.model(
            backbone=self.backbone, views=views, meta=meta, cameras=cameras, resize_transform=resize_transform
        )
        return fused_poses, plane_poses, proposal_centers, input_heatmaps

dummy_inputs = {
    "views": torch.randn(1, 5, 3, 512, 960).to(device),    # Example for multi-view inputs
    "meta": {"seq": ["customized_sequence"]},  # Example metadata
    "cameras": camera,            # Example camera parameters
    "resize_transform": torch.randn(2, 3).to(device) # Example transformation
}
torch.onnx.export(combined_model, (dummy_inputs["views"], dummy_inputs["meta"], 
     dummy_inputs["cameras"], dummy_inputs["resize_transform"]), onnx_filename, 
                  export_params=True,   # Export trained parameters
                  opset_version=16,
                  verbose=False,#
                  do_constant_folding=True,  # Optimization for constant folding
                  input_names=["views", "meta", "cameras", "resize_transform"],
                output_names=["fused_poses", "plane_poses", "proposal_centers", "input_heatmaps"], # Name for the output node
                  dynamic_axes={'input' : {0 : 'batch_size'},    # Dynamic batch size
                                'output' : {0 : 'batch_size'}})