i’m trying to iterate through a dataloader created from a pytorch dataset. when I run the loop it gives me "index 1 is out of bounds for dimension 0 with size 1 error ". below is my dataset class
class traindataset(Dataset):
class traindataset(Dataset):
    def __init__(self,data,train_end_idx,augmentation=None):
        '''
        data: data is a pandas dataframe generated from csv file where it has columns-> [name,labels,col 1,col2,...,col784]. shape of data->(10000, 786)
        
        
        '''
        self.data=data
        self.augmentation=augmentation
        self.train_end=train_end_idx
        self.target=self.data.iloc[:self.train_end,1].values
        
    def __len__(self):
        return len(self.target);
    def __getitem__(self,idx):
        target=self.target
        image=self.data.iloc[:self.train_end,2:].values
        if self.augmentation is not None:
            image = self.augmentation(image)
        return torch.tensor(target[idx]),image[idx]


below is my augmenataion and dataloader generator
torchvision_transform = transforms.Compose([
    np.uint8,
    transforms.ToPILImage(),
    transforms.Resize((28,28)),
    transforms.RandomRotation([45,135]),
    transforms.ToTensor()
    ])

below is the loop I’m running, where I’m getting the mentioned error
for label,image in trainloader:
    print(label,train)

below is the complete error I have received.
IndexError                                Traceback (most recent call last)
/tmp/ipykernel_41/1540740000.py in <module>
----> 1 for label,image in trainloader:
      2     print(label,train)

/opt/conda/lib/python3.7/site-packages/torch/utils/data/dataloader.py in __next__(self)
    519             if self._sampler_iter is None:
    520                 self._reset()
--> 521             data = self._next_data()
    522             self._num_yielded += 1
    523             if self._dataset_kind == _DatasetKind.Iterable and \

/opt/conda/lib/python3.7/site-packages/torch/utils/data/dataloader.py in _next_data(self)
    559     def _next_data(self):
    560         index = self._next_index()  # may raise StopIteration
--> 561         data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    562         if self._pin_memory:
    563             data = _utils.pin_memory.pin_memory(data)

/opt/conda/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py in fetch(self, possibly_batched_index)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---> 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

/opt/conda/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py in <listcomp>(.0)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---> 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

/tmp/ipykernel_41/1814544681.py in __getitem__(self, idx)
     18         if self.augmentation is not None:
     19             image = self.augmentation(image)
---> 20         return torch.tensor(target[idx]),image[idx]

IndexError: index 1 is out of bounds for dimension 0 with size 1

note: code works fine without augmentation.