Support num_samples < batch_size

[csala]

The current implementation ends up in an infinite loop if the num_samples is lower than the batch_size.

This should be prevented by either always sampling at least one batch, or making the batch_size equal to the num_samples if the later is lower.