Thanks for your response. I have changed the initialization code to:
sample = CustomSampler(data_source, 3)
but error was the same 
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
 in ()
----> 1 sample = CustomSampler(data_source, 3)
TypeError: CustomSampler() takes 1 positional argument but 2 were given