Solved by ptrblck in post #2 
              
              
                It seems you are using a custom collate_fn with an if-condition: 
if random.randint(2)

The random.randint method takes two arguments: the lower bound a and the upper bound b. 
If you would like to sample in [0, 2], you can use random.randint(0, 2) instead. 
Note that the Python randint method will …