Solved by ptrblck in post #2 
              
              
                I’m not exactly sure, what your use case is, but from the code perspective, it should work, if you use 
remaining_sticks = remaining_sticks * (1-km[:,i])

instead of the inplace operation. 
Additionally, you could also initialize sticks as a list and use torch.cat to create a tensor afterwards.