Add device to pre-processing

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

nemotron_page_elements_v3/model.py

@@ -145,6 +145,7 @@ class YoloXWrapper(nn.Module):
         """
         if not isinstance(image, torch.Tensor):
             image = torch.from_numpy(image)
+        image = image.to(self.device)
         image = image.permute(2, 0, 1)  # [H, W, 3] -> [3, H, W]
         image = resize_pad(image, self.img_size)
         return image.float()