vjawa_ensure_device_side_preprocessing (#8)

- Add device to pre-processing (0e5c49d842cf2fb0ccadd87fe65a6eff06d8da77)


Co-authored-by: Vibhu Jawa <VibhuJawa@users.noreply.huggingface.co>

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()