In all there are 8 labels,my custom dataset is as follows
classes=('m_as01','f_as01','m_af01','f_af01','m_ca01','f_ca01','m_la01','f_la01')
idx_to_class = {i:j for i, j in enumerate(classes)}
class_to_idx = {value:key for key,value in idx_to_class.items()}

class BlenderPoseDataset(Dataset):
    
    def __init__(self,paths,batch_size=16):
        self.img_dir='blender_data_1/images_blender/*'
        self.pose_files=paths
        self.batch_size=batch_size
        self.transforms=None
        self.image_shape=(224,224)
        self.data=[]
        img_list=glob.glob(self.img_dir)
        for img_path in img_list:
            for file in self.pose_files:
                self.data.append([img_path,file])
    def __len__(self):
        return len(self.data)
    
    def image_process(self,rgb_img):
        rgb_img=np.transpose(rgb_img.astype('float32'),(2,0,1))/255.0
        return rgb_img
    
    
    def __getitem__(self,idx):
        if self.transforms is not None:
            self.image=self.transform(self.image)
        img_path,label_path=self.data[idx]
        img=cv2.imread(img_path)
        img=cv2.resize(img,self.image_shape)[:,:,::-1].astype(np.float32)
        img=self.image_process(img)
        img=torch.from_numpy(img).float()
        with open(label_path,'r') as f:
          label=f.readline()
        label=class_to_idx[label]
        return img,label

output when iterating through train loader
Batch of image has shape: torch.Size([128, 3, 224, 224])
Batch of pose has shape: torch.Size([128])