mirrors
/
DCL
mirror of https://github.com/JDAI-CV/DCL.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update dataset_DCL.py

master
SaiYiKi 2019-07-23 23:07:44 +08:00 committed by GitHub
parent 3b5b5e95ac
commit d9ddacda59
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
utils/dataset_DCL.py
View File

@ -162,10 +162,10 @@ def collate_fn4backbone(batch):
img_name = [] img_name = []
for sample in batch: for sample in batch:
imgs.append(sample[0]) imgs.append(sample[0])
if len(sample) == 5: if len(sample) == 7:
label.append(sample[1])
else:
label.append(sample[2]) label.append(sample[2])
else:
label.append(sample[1])
img_name.append(sample[-1]) img_name.append(sample[-1])
return torch.stack(imgs, 0), label, img_name return torch.stack(imgs, 0), label, img_name