Ryan Chesler commited on
Commit
5b3a642
·
1 Parent(s): 4ba3706

Fix undefined idx variable in quad_non_maximal_suppression callThe for loop was removed but the [idx].unsqueeze(0) indexing wasleft behind, causing a NameError. Pass the full batch tensorsdirectly since the CUDA kernel now supports batch mode natively

Browse files
nemotron-ocr/src/nemotron_ocr/inference/pipeline.py CHANGED
@@ -182,8 +182,8 @@ class NemotronOCR:
182
  e2e_det_coords = rrect_to_quads(det_rboxes.float(), DETECTOR_DOWNSAMPLE)
183
 
184
  quads, confidence, region_counts = quad_non_maximal_suppression(
185
- e2e_det_coords[idx].unsqueeze(0),
186
- e2e_det_conf[idx].unsqueeze(0),
187
  prob_threshold=NMS_PROB_THRESHOLD,
188
  iou_threshold=NMS_IOU_THRESHOLD,
189
  kernel_height=2,
 
182
  e2e_det_coords = rrect_to_quads(det_rboxes.float(), DETECTOR_DOWNSAMPLE)
183
 
184
  quads, confidence, region_counts = quad_non_maximal_suppression(
185
+ e2e_det_coords,
186
+ e2e_det_conf,
187
  prob_threshold=NMS_PROB_THRESHOLD,
188
  iou_threshold=NMS_IOU_THRESHOLD,
189
  kernel_height=2,