Hi I have been writing a code in Pytorch for my custom sampler. But when I run it, it generates an error.
CODE
import torch
from torch.utils.data.sampler import Sampler

def CustomSampler(Sampler):
    def __init__(self, data_source, batch_size, replacement=False):
        self.data_source = data_source
        self.ids = torch.FloatTensor(data_source.keys())
        self.numSmall = torch.FloatTensor(data_source.values())
        self.normal_numSmall = self.numSmall / self.numSmall.sum()

        self.num_samples = numSmall.numel()
        self.batch_size = batch_size

    def __iter__(self):
        return self.ids[torch.multinomial(self.normal_small, self.batch_size, replacement=False).numpy().tolist()]

    def __len__(self):
        return self.batch_size

id = [11,12,156,256,36,26,21,25]
small = [13,7,6,8,12,15,3,20]

data_source = dict(zip(id, small))
sample = CustomSampler(data_source, batch_size = 3)

Error
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-9-d53f4a02f4a7> in <module>()
----> 1 sample = CustomSampler(data_source, batch_size = 3)

TypeError: CustomSampler() got an unexpected keyword argument 'batch_size'

Could you please tell me what is the source of error? How can I settle this issue?