Fix float for imagenet rightfit

Author: Amar3tto

sdks/python/apache_beam/examples/inference/pytorch_imagenet_rightfit.py

@@ -88,12 +88,14 @@ def decode_and_preprocess(image_bytes: bytes, size: int = 224) -> torch.Tensor:
 
     # To tensor [0..1]
     import numpy as np
+    mean = np.array(IMAGENET_MEAN, dtype=np.float32)
+    std = np.array(IMAGENET_STD, dtype=np.float32)
     arr = np.asarray(img).astype("float32") / 255.0  # H,W,3
     # Normalize
-    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
+    arr = (arr - mean) / std
     # HWC -> CHW
-    arr = np.transpose(arr, (2, 0, 1))
-    return torch.from_numpy(arr)  # float32, shape (3,224,224)
+    arr = np.transpose(arr, (2, 0, 1)).astype("float32")
+    return torch.from_numpy(arr).float()  # float32, shape (3,224,224)
 
 
 class RateLimitDoFn(beam.DoFn):
@@ -174,9 +176,28 @@ def __init__(self, top_k: int, model_name: str):
 
   def process(self, kv: Tuple[str, PredictionResult]):
     image_id, pred = kv
-    logits = pred.inference[
-        "logits"]  # torch.Tensor [B, num_classes] or [num_classes]
-    if isinstance(logits, torch.Tensor) and logits.ndim == 1:
+
+    # pred can be PredictionResult OR raw inference object.
+    inference_obj = pred.inference if hasattr(pred, "inference") else pred
+
+    # inference_obj can be dict {'logits': tensor} OR tensor directly.
+    if isinstance(inference_obj, dict):
+      logits = inference_obj.get("logits", None)
+      if logits is None:
+        # fallback: try first value if dict shape differs
+        try:
+          logits = next(iter(inference_obj.values()))
+        except Exception:
+          logits = None
+    else:
+      logits = inference_obj
+
+    if not isinstance(logits, torch.Tensor):
+      logging.warning("Unexpected logits type for %s: %s", image_id, type(logits))
+      return
+
+    # Ensure shape [1, C]
+    if logits.ndim == 1:
       logits = logits.unsqueeze(0)
 
     probs = F.softmax(logits, dim=-1)  # [B, C]
@@ -480,7 +501,7 @@ def run(
 
   to_infer = (
       preprocessed
-      | 'ToKeyedTensor' >> beam.Map(lambda kv: (kv[0], kv[1]["tensor"])))
+      | 'ToKeyedTensor' >> beam.Map(lambda kv: (kv[0], kv[1]["tensor"].float())))
 
   predictions = (
       to_infer