Is there a function working as tf.segment_max() in pytorch?

1

As shown in the title