Repository: pytorch
Issue: Setting CPU affinity doesn't work for torch.softmax when using a thread
Labels: module: cpu,triaged

### 🐛 Describe the bug

This bug is about the cpu affinity.
When I use torch.softmax in a thread of a child process, other cpus are used even though I set cpu affinity.
But if I comment out torch.sofmax in the function, only cpus that I set cpu affinity are used.
I think it is a bug of torch.softmax.
I tested torch.sigmoid, and torch.tanh but they don't make this situation.

```python
import os

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"

import time
import psutil
import torch
import torch.multiprocessing as mp
from torch import nn

from threading import Thread


def f(model):
    while True:
        x = torch.randn(2, 256)
        with torch.no_grad():
            y = model(x)
            p = torch.softmax(y, dim=-1)  # This line is the source of the problem


def run_worker(cpu_id):
    p = psutil.Process()
    p.cpu_affinity([cpu_id])

    print(f"Process {p.pid} is running on CPU {cpu_id}")
    model = nn.Linear(256, 256)
    th1 = Thread(target=f, args=(model,))
    th1.start()

    while True:
        time.sleep(1e-2)


if __name__ == "__main__":
    mp.set_start_method("spawn")

    torch.set_num_threads(1)

    procs = list()
    for i in range(2):
        p = mp.Process(target=run_worker, args=(i,))
        p.start()
        procs.append(p)

    for p in procs:
        p.join()
```
With out softmax:
![without_softmax](https://github.com/user-attachments/assets/26ae7a6a-3ff0-456d-ba1c-2e800caf1d35)

With softmax:
![with_softmax](https://github.com/user-attachments/assets/7599d958-02ad-43d2-9d9a-b7f995e2ebca)

If I change softmax as below, it resolves the problem.
```python
            # p = torch.softmax(y, dim=-1)
            y = torch.exp(y)
            p = y / y.sum(dim=-1, keepdim=True)
```

But still, if I sample using torch.distributions.Categorical(p).sample() cause the sample problem again.
```python
            # p = torch.softmax(y, dim=-1)
            y = torch.exp(y)
            p = y / y.sum(dim=-1, keepdim=True)
            a = torch.distributions.Categorical(p).sample()
```


### Versions

[pip3] numpy==2.1.0
[pip3] torch==2.4.0
[pip3] torchaudio==2.4.0
[pip3] torchvision==0.19.0
[conda] blas                      2.116                       mkl    conda-forge
[conda] blas-devel                3.9.0            16_linux64_mkl    conda-forge
[conda] cpuonly                   2.0                           0    pytorch
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            16_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     2.1.0           py311h71ddf71_1    conda-forge
[conda] pytorch                   2.4.0              py3.11_cpu_0    pytorch
[conda] pytorch-mutex             1.0                         cpu    pytorch
[conda] torchaudio                2.4.0                 py311_cpu    pytorch
[conda] torchvision               0.19.0                py311_cpu    pytorch

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10