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