Przeglądaj źródła

fix: ensure batch_size is at least 1 when sorted_images is empty in Unimernet.py

myhloli 3 miesięcy temu
rodzic
commit
a636b34324
1 zmienionych plików z 1 dodań i 1 usunięć
  1. 1 1
      mineru/model/mfr/unimernet/Unimernet.py

+ 1 - 1
mineru/model/mfr/unimernet/Unimernet.py

@@ -106,7 +106,7 @@ class UnimernetModel(object):
         dataset = MathDataset(sorted_images, transform=self.model.transform)
 
         # 如果batch_size> len(sorted_images),则设置为不超过len(sorted_images)的2的幂
-        batch_size = min(batch_size, 2 ** (len(sorted_images).bit_length() - 1))
+        batch_size = min(batch_size, max(1, 2 ** (len(sorted_images).bit_length() - 1))) if sorted_images else 1
 
         dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=0)