vjawa_move_preprocessing_to_device (#8)

- move pre processing to device (0f79b282d3df71101f492bc42acd6c49beb2957a)


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

nemotron_graphic_elements_v1/model.py

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