@ptrblck thank you for patiently and quickly replying, I have made the below changes in my code but it still has the same issue.
class traindataset(Dataset):
    def __init__(self,data,train_end_idx,augmentation = None):
        '''
        data: pandas dataframe
        
        '''
        self.data=data
        self.augmentation=augmentation
        self.train_end=train_end_idx
        self.target=self.data.iloc[:self.train_end,1].values
        self.image=self.data.iloc[:self.train_end,2:].values #contains full data
        
    def __len__(self):
        return len(self.target);
    def __getitem__(self,idx):
        self.target=self.target
        image=self.image[idx].reshape(1,784) #only takes the selected index
        if self.augmentation is not None:
            image = self.augmentation(image)
        
        return torch.tensor(self.target[idx]),torch.tensor(image)

changes I have done are:

in init I put my whole dataset rather than in getitem()->image=self.image[idx].reshape(1,784)
in get item I’m passing only one image by specifying its index
->image=self.image[idx].reshape(1,784)