I am trying to write a model in pytorch. I have 3 classes A,B,C. Each class is a list of tensors. I want to take samples from these classes. Each class has also a reserve class that when the class is empty, it take sample from it. For example when class A becomes empty, my model take sample from ...