vjawa_move_preprocessing_to_device

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