diff --git a/make_patch.py b/make_patch.py
index 58e40cf..065fbea 100755
--- a/make_patch.py
+++ b/make_patch.py
@@ -90,7 +90,7 @@
 idx = np.arange(50000)
 np.random.shuffle(idx)
 training_idx = idx[:train_size]
-test_idx = idx[train_size:test_size]
+test_idx = idx[train_size:(train_size + test_size)]
 
 train_loader = torch.utils.data.DataLoader(
     dset.ImageFolder('./imagenetdata/val', transforms.Compose([