From 88136aa97efc55b2dd2964574346893f54ba81bf Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Tue, 12 Mar 2019 10:58:56 -0700 Subject: [PATCH] fix behavior of ConcatDataset w/ negative indices (#15756) Summary: Currently, when you pass a negative index to a `Dataset` created with `ConcatDataset`, it simply passes that index to the first dataset in the list. So if, for example, we took `concatenated_dataset[-1]`, this will give us the last entry of the *first* dataset, rather than the last entry of the *last* dataset, as we would expect. This is a simple fix to support the expected behavior for negative indices. Pull Request resolved: https://github.com/pytorch/pytorch/pull/15756 Reviewed By: ezyang Differential Revision: D14081811 Pulled By: fmassa fbshipit-source-id: a7783fd3fd9e1a8c00fd076c4978ca39ad5a8a2a --- torch/utils/data/dataset.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/torch/utils/data/dataset.py b/torch/utils/data/dataset.py index 0202ca25f7ad..bb688ce6bf22 100644 --- a/torch/utils/data/dataset.py +++ b/torch/utils/data/dataset.py @@ -73,6 +73,10 @@ def __len__(self): return self.cumulative_sizes[-1] def __getitem__(self, idx): + if idx < 0: + if -idx > len(self): + raise ValueError("absolute value of index should not exceed dataset length") + idx = len(self) + idx dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) if dataset_idx == 0: sample_idx = idx