Skip to content

Commit 1d654cd

Browse files
committed
Fix inference
1 parent 7dd82b3 commit 1d654cd

1 file changed

Lines changed: 4 additions & 1 deletion

File tree

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,10 @@ def __init__(self, blip_name: str, clip_name: str):
111111

112112
def process(self, kv: Tuple[str, PredictionResult]):
113113
image_id, pred = kv
114-
inf = pred.inference or {}
114+
if hasattr(pred, "inference"):
115+
inf = pred.inference or {}
116+
else:
117+
inf = pred
115118
# Expected inference fields from CLIP handler:
116119
# best_caption, best_score, candidates, scores, blip_ms, clip_ms, total_ms
117120
best_caption = inf.get("best_caption", "")

0 commit comments

Comments
 (0)