Fix AutoBatch when working with RT-DETR models (#18912)

This commit is contained in:
Laughing 2025-01-27 18:04:03 +08:00 committed by GitHub
parent 305a298ae2
commit 30a2de164b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 2 additions and 2 deletions

View file

@ -172,7 +172,7 @@ def get_cdn_group(
bounding boxes, attention mask and meta information for denoising. If not in training mode or 'num_dn'
is less than or equal to 0, the function returns None for all elements in the tuple.
"""
if (not training) or num_dn <= 0:
if (not training) or num_dn <= 0 or batch is None:
return None, None, None, None
gt_groups = batch["gt_groups"]
total_num = sum(gt_groups)