Repository: pytorch
Issue: Driver Allocated Memory grows unrestricted when using torch.unique on MPS device
Labels: module: memory usage,triaged,module: mps

### 🐛 Describe the bug

When using `torch.unique`  in a loop on the MPS backend, the memory allocated by the driver grows unrestricted. In my real application that leads to an  `RuntimeError: MPS backend out of memory (MPS allocated: 24.00 MB, other allocations: 36.24 GB, max allowed: 36.27 GB)` error late in the training.

I created this minimal example with the same behaviour.

```python
import torch
import gc


def test_operations(iterations: int, shape: tuple[int, int]) -> None:
    print(f"PyTorch version: {torch.__version__}")
    # Test 1: torch.unique
    print("\nTest 1: torch.unique")
    x = torch.randint(0, 2, shape, device="mps")
    for i in range(iterations):
        y = torch.unique(x)
        del y

        # Empty cache and collect garbage to make sure
        torch.mps.empty_cache()
        gc.collect()

        if i % 10 == 0:
            print(
                f"Iter {i}: Driver Allocated Memory: {torch.mps.driver_allocated_memory() / (1024**2):.2f}MB, Current Allocated Memory: {torch.mps.current_allocated_memory() / (1024**2):.2f}MB"
            )

    # Test 2: torch.sort (comparison)
    print("\nTest 2: torch.sort")
    for i in range(iterations):
        y = torch.sort(x)[0]
        del y
        # Empty cache and collect garbage to make sure
        torch.mps.empty_cache()
        gc.collect()

        if i % 10 == 0:
            print(
                f"Iter {i}: Driver memory: {torch.mps.driver_allocated_memory() / (1024**2):.2f}MB, Current memory: {torch.mps.current_allocated_memory() / (1024**2):.2f}MB"
            )


test_operations(iterations=100, shape=(2000, 10))
```

Results in 

```
PyTorch version: 2.5.1

Test 1: torch.unique
Iter 0: Driver Allocated Memory: 18.73MB, Current Allocated Memory: 0.15MB
Iter 10: Driver Allocated Memory: 98.73MB, Current Allocated Memory: 0.15MB
Iter 20: Driver Allocated Memory: 178.73MB, Current Allocated Memory: 0.15MB
Iter 30: Driver Allocated Memory: 258.73MB, Current Allocated Memory: 0.15MB
Iter 40: Driver Allocated Memory: 338.73MB, Current Allocated Memory: 0.15MB
Iter 50: Driver Allocated Memory: 418.73MB, Current Allocated Memory: 0.15MB
Iter 60: Driver Allocated Memory: 578.73MB, Current Allocated Memory: 0.15MB
Iter 70: Driver Allocated Memory: 738.72MB, Current Allocated Memory: 0.15MB
Iter 80: Driver Allocated Memory: 898.72MB, Current Allocated Memory: 0.15MB
Iter 90: Driver Allocated Memory: 1058.72MB, Current Allocated Memory: 0.15MB

Test 2: torch.sort
Iter 0: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 10: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 20: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 30: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 40: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 50: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 60: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 70: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 80: Driver memory: 1202.72MB, Current memory: 0.15MB
Iter 90: Driver memory: 1202.72MB, Current memory: 0.15MB
```

Showing the increase in the driver allocated memory when using `torch.unique` but not when using another function like `torch.sort`. I just used `torch.sort` as comparison here.

Is this behaviour expected?

### Versions

Collecting environment information...
PyTorch version: 2.5.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.2 (arm64)
GCC version: Could not collect
Clang version: 16.0.0 (clang-1600.0.26.6)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.11 (main, Dec  3 2024, 17:20:40) [Clang 16.0.0 (clang-1600.0.26.4)] (64-bit runtime)
Python platform: macOS-15.2-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2 Pro

Versions of relevant libraries:
[pip3] torch==2.5.1
[conda] Could not collect

cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen