vjawa_move_preprocessing_to_device (#5)

- Add pre processing to device (917dc98097032dbc722ec48a2ba5d4dce642d192)


Co-authored-by: Vibhu Jawa <[email protected]>

nemotron_table_structure_v1/model.py

@@ -146,6 +146,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()