Dataloader taking too long on google colab?

This is all on google colab.
I don’t think there should be anything wrong with the code. there is no loop in getItem of the dataset. But when I try running the code. First, it gets stuck for a very long time on colab on GPU; second, sometimes it just works for one iteration .
__First I try it on CPU and it worked just fine. As soon as I switched to GPU everything went wrong, and it doesn’t work on CPU either [might not be the problem I am sure but worth mentioning __
Here is the code for dataset :

    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        with  open(os.path.join(self.root_dir,'tid2013_Dis_Ref_directory.txt'),'r') as f  :
          self.disRefs =[[s.replace('\\',os.path.sep).strip()  for s in line.split(' # ')] for line in f.readlines()]
    def __len__(self):
        return len(self.disRefs)
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        img_name = os.path.join(self.root_dir,self.disRefs[idx][0])
        image = Image.fromarray(io.imread(img_name))
        if self.transform:
            image = self.transform(image)
        sample = [ self.disRefs[idx][0],self.disRefs[idx][1] ,image]#.unsqueeze(0)] 
        return sample
    def RefTransform (self, path ):
      img_name = os.path.join(self.root_dir,path)
      img = Image.fromarray(io.imread(img_name))
      if self.transform:
          img = self.transform(img)
      img = img.unsqueeze(0)
      return img
root_dir = '/content/gdrive/My Drive/tid2013'
dataSet = IQADataset(root_dir , transform_)

and here is the main for loop :

batch_size_= 1
dataloader = DataLoader(dataSet, batch_size=batch_size_,shuffle=False, num_workers= 2)
PerceptualLosses ={}
device_ = torch.device("cuda" if torch.cuda.is_available() else "cpu")
count = 0
start = datetime.utcnow()
for dis_dirs,ref_dirs,images in dataloader:  
  if count >=120 or count == 0:
    count = 0 
    img_name = os.path.join(root_dir,
    refImg = Image.fromarray(io.imread(img_name))
    refImg = transform_(refImg)
    refImg = refImg.unsqueeze(0)
    refImg =
    print (f'************{ref_dirs[0]}')
    ref_feats = model(refImg)
    ref_feats =  deepcopy( modelLayersOutput)
    ref_feats = [np.repeat(np.expand_dims(layerOutput,0), batch_size_, axis=0) for layerOutput in ref_feats]
    images =
    res = model(images)
    res = getMSE( ref_feats , modelLayersOutput )
    print(str(datetime.utcnow() - start ))
    for i in range (len (range(batch_size_))):
      end = datetime.utcnow()
      print (f'{index} #{dis_dirs[i]}')
      PerceptualLosses[dis_dirs[i]] = [ref_dirs[i] ,[re[i] for re in res ]]

Can you add a timeit function in your dataloader’s __getitem__() function just to see which part takes the longest time?

just to mention sth . dataloader iterates all the way through the end ; that’s why its stuck in the for loop . Am I missing sth here ?

I think @azhanmohammed suggestion was to profile the data loading and to check how long the loading and processing of a single batch takes in the Colab environment. You are right that the DataLoader loop will iterate the entire dataset, so you could profile a few iterations and then use a break statement to end the loop.