From 35abefd038c64b5bd72b7fe86ba521055e814259 Mon Sep 17 00:00:00 2001 From: James Betker Date: Sat, 1 Jan 2022 10:31:03 -0700 Subject: [PATCH] More fix --- codes/data/zero_pad_dict_collate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/data/zero_pad_dict_collate.py b/codes/data/zero_pad_dict_collate.py index 5423f63a..09b3a6fa 100644 --- 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 \ No newline at end of file