how to load few samples per class about medmnist HOT 4 CLOSED

data_flag = 'chestmnist'
download = True
classes=2
NUM_EPOCHS = 3
BATCH_SIZE = 128
lr = 0.001

info = INFO[data_flag]
task = info['task']
n_channels = info['n_channels']
n_classes = len(info['label'])

DataClass = getattr(medmnist, info['python_class'])
# preprocessing
data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5], std=[.5])
])

# load the data
train_dataset = DataClass(split='train', transform=data_transform, download=download,)
test_dataset = DataClass(split='test', transform=data_transform, download=download)

pil_dataset = DataClass(split='train', download=download)

# encapsulate data into dataloader form
train_loader = data.DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
train_loader_at_eval = data.DataLoader(dataset=train_dataset, batch_size=2*BATCH_SIZE, shuffle=False)
test_loader = data.DataLoader(dataset=test_dataset, batch_size=2*BATCH_SIZE, shuffle=False)

import numpy as np

# Load the .npz file
data = np.load('chestmnist.npz')

# Get the names of the arrays in the file
array_names = data.files

# Create an empty dictionary to store the arrays
arrays_dict = {}

# Iterate over the array names and store the arrays in the dictionary
for name in array_names:
    arrays_dict[name] = data[name]

['train_images', 'val_images', 'test_images', 'train_labels', 'val_labels', 'test_labels']