updates
This commit is contained in:
parent
9b84885775
commit
4b56a370e6
|
@ -499,7 +499,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
|||
|
||||
@staticmethod
|
||||
def collate_fn(batch):
|
||||
img, label, path, shapes = list(zip(*batch)) # transposed
|
||||
img, label, path, shapes = zip(*batch) # transposed
|
||||
for i, l in enumerate(label):
|
||||
l[:, 0] = i # add target image index for build_targets()
|
||||
return torch.stack(img, 0), torch.cat(label, 0), path, shapes
|
||||
|
|
Loading…
Reference in New Issue