torch_take_along_dim function

Selects values from input at the 1-dimensional indices from indices along the given dim.