Using RandWeightedCropd with correct batch size

[sunburnt-fish]

I'm trying to use RandWeightedCropd to generate samples from a volume of data. For each volume I want to generate 50 samples but I cannot fit all of this onto my GPU so I need a smaller batch size. However, when I apply this transform and set the batch size in my data loader, it always returns a batch size of num_samples that I set for the random sampling and not the smaller batch_size I set for my data loader. How can I do this correctly? A minimal example is below to show my problem:

import numpy as np
from monai.data import DataLoader, Dataset
from monai.transforms import RandWeightedCropd

if __name__ == '__main__':
    volume_size = (1, 256, 256, 256)
    data = {
        'volume': np.random.random(volume_size),
        'weights': np.random.random(volume_size)
    }
    keys = data.keys()
    transforms = RandWeightedCropd(keys=keys, w_key='weights', spatial_size=(128, 128, 128), num_samples=50)
    dataset = Dataset(data=[data], transform=transforms)
    dataloader = DataLoader(dataset=dataset, batch_size=10)

    for batch in dataloader:
        batch_volume = batch['volume'].numpy()
        batch_weights = batch['weights'].numpy()
        print(batch_volume.shape)
        print(batch_weights.shape)
    exit(0)

I would expect this to give me a print out of

(10, 1, 128, 128, 128)
(10, 1, 128, 128, 128)

but instead I get

(50, 1, 128, 128, 128)
(50, 1, 128, 128, 128)

[KumoLiu]

Hi @sunburnt-fish, you can refer to this discussion.
#6242

Hope it helps, thanks!