a-One-Fan/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Pull Requests Activity

More fix

Browse Source
This commit is contained in:
James Betker 2022-01-01 10:31:03 -07:00
parent d5a5111890
commit 35abefd038
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/data/zero_pad_dict_collate.py
View File

@ -39,7 +39,7 @@ class ZeroPadDictCollate():
if len(first_dict[key].shape) > 0: if len(first_dict[key].shape) > 0:
collated[key] = self.collate_tensors(batch, key) collated[key] = self.collate_tensors(batch, key)
else: else:
collated[key] = torch.stack(batch[key]) collated[key] = torch.stack([b[key] for b in batch])
else: else:
collated[key] = self.collate_into_list(batch, key) collated[key] = self.collate_into_list(batch, key)
return collated return collated