- High: It blocks me to complete my task.

for the example ：Image Classification Batch Inference with PyTorch — Ray 2.9.0

```
import numpy as np
from typing import Any, Dict
class ResnetModel:
def __init__(self):
self.weights = ResNet152_Weights.IMAGENET1K_V1
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = models.resnet152(weights=self.weights).to(self.device)
self.model.eval()
def __call__(self, batch: Dict[str, np.ndarray]):
# Convert the numpy array of images into a PyTorch tensor.
# Move the tensor batch to GPU if available.
torch_batch = torch.from_numpy(batch["transformed_image"]).to(self.device)
with torch.inference_mode():
prediction = self.model(torch_batch)
predicted_classes = prediction.argmax(dim=1).detach().cpu()
predicted_labels = [
self.weights.meta["categories"][i] for i in predicted_classes
]
return {
"predicted_label": predicted_labels,
"original_image": batch["original_image"],
}
def preprocess_image(row: Dict[str, np.ndarray]):
return {
"original_image": row["image"],
"transformed_image": transform(row["image"]),
}
transformed_ds = ds.map(preprocess_image)
predictions = transformed_ds.map_batches(
ResnetModel,
concurrency=4, # Use 4 GPUs. Change this number based on the number of GPUs in your cluster.
num_gpus=1, # Specify 1 GPU per model replica.
batch_size=720, # Use the largest batch size that can fit on our GPUs
)
```

the predictions result is batch result,

How to match the inference result and the dataset id ,

for item in ds :

item[id]

item[predicted_label]