More fix

2022-01-01 10:31:03 -07:00 · 2022-01-01 10:31:03 -07:00 · 35abefd038
commit 35abefd038
parent d5a5111890
1 changed files with 1 additions and 1 deletions
--- a/codes/data/zero_pad_dict_collate.py
+++ b/codes/data/zero_pad_dict_collate.py
@ -39,7 +39,7 @@ class ZeroPadDictCollate():
                if len(first_dict[key].shape) > 0:
                    collated[key] = self.collate_tensors(batch, key)
                else:
-                    collated[key] = torch.stack(batch[key])
+                    collated[key] = torch.stack([b[key] for b in batch])
            else:
                collated[key] = self.collate_into_list(batch, key)
        return collated