I need to use an equivalent of  tf.dynamic_partition  in PyTorch. Is there anything with similar functionality in PyTorch or other library or is there a simple and clever way to code it for PyTorch and work fast? The same for tf.dynamic_stitch. Thanks!