Replace argsort to reduce memory footprint
@dkurt 
Thank you for preparing this PR.
Let me test this and review it.
Need to verify on multiple samples.  @tango4j, any updates? This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. @tango4j , friendly reminder. Maybe I can run some tests on my side? This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. @tango4j, reminder This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. This PR was closed because it has been inactive for 7 days since being marked as stale.
# What does this PR do ?

Reduce memory footprint by replacing torch.argsort. With my local sample, diarization pipeline works on 12GB. Without the patch, it failed with an error `torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.60 GiB. GPU 0 has a total capacty of 23.64 GiB of which 3.21 GiB is free. Including non-PyTorch memory, this process has 20.35 GiB memory in use. Of the allocated memory 16.07 GiB is allocated by PyTorch, and 665.59 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF`

**Collection**: [asr]

# Changelog 
Replace torch.argsort to reduce memory usage

# Usage
Running MSDD on long audio sample

# Jenkins CI
To run Jenkins, a NeMo User with write access must comment `jenkins` on the PR.

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [x] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)
