Imagenet Dataloading Correct?

I have attached dataloader that I have written for Imagenet. I have saved the paths to images with their respective classname in .npy file and load it from there. Hierarchy of file is same as provided in Kaggle dataset for Imagenet. I am not sure if it is correct. It seems to return image tensor and label fine, but model accuracy is very low - so I was wondering if I did any mistake in this part. Any feedback will be highly appreciated! Thank you!

# Custom Dataloader
class CustomDataset(Dataset):
    def __init__(self,train=True,transform=None):
        self.data = []
        self.transform = transform 
        # Class to Label Encoding Loading
        class_to_label_loading = np.load("~/model_implementation/AlexNet/class_to_label.npy",allow_pickle=True)
        self.class_to_label = class_to_label_loading.item()
        if not train:
            self.data = np.load("~/model_implementation/AlexNet/val_data.npy")
        else:
            self.data = np.load("~/model_implementation/AlexNet/train_data.npy")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path, class_name = self.data[idx]
        img = Image.open(img_path)
        rgb_img = img.convert('RGB')
        label = self.class_to_label[class_name]
        if self.transform is not None:
            img_tensor = self.transform(rgb_img)
        label_tensor = torch.tensor([label])
        return img_tensor, label_tensor

and transform is pretty standard one for Imagenet dataset like this:

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

Based on your code it seems as if you’ve stored the paths in the numpy array only. I cannot see any obvious errors, but since the data structure is equal to the “Kaggle dataset” (assuming this means the samples for each class are moved to their respective class folder) you could compare the model training with the standard ImageFolder dataset.