when counts[k] is zero, an stack error occurs.

[Yongha Kwon]

In sample function, when counts[k] is zero, x_k = torch.stack([d_k.sample() for _ in range(int(counts[k]))]) is error.

So, I think need add this.

if counts[k] == 0:
  continue

[Lucas Deecke]

Fixed in 156031e — thank you for noting this @Danden1! 👍