Image в Tensor data

Библиотека: pytorch_java_only-1.6.0

Я хочу конвертировать BufferedImage/File в Tensor data, есть ли какой-то метод/библиотека для этого?

Python код:

image = Image.open(image_path)
image = image.convert('RGB') 
transform = Compose([ToTensor()])
image = transform(image)
image = image.view(1, 3, 64, 64).cuda()
output = my_model(image)
output = output.view(-1, self.quantity)
output = nn.functional.softmax(output, dim=1)
output = torch.argmax(output, dim=1)
output = output.view(-1, self.size)[0]

Я посмотрел java-demo

И вот что меня интересует:

Module mod = Module.load("demo-model.pt1");
Tensor data =
Tensor.fromBlob(
    new int[]  {1, 2, 3, 4, 5, 6}, // **data array**
    new long[] {2, 3} // **shape array**
    );
IValue result = mod.forward(IValue.from(data), IValue.from(3.0));
Tensor output = result.toTensor();
System.out.println("shape: " + Arrays.toString(output.shape()));
System.out.println("data: " + Arrays.toString(output.getDataAsFloatArray()));

С значениями shape все ясно, это HEIGHT, WIDTH, CHANNELS картинки.

Но что насчёт data? Как я могу конвертировать картинку в правильный массив? Возможно мне стоит использовать javacv/nd4j или что-то в этом роде?

Спасибо за ответ!