Repository: pytorch
Issue: Environment gating of CUDA_VISIBLE_DEVICES returns a CUDA initialization error.
Labels: high priority,module: cuda,triaged,module: regression

### 🐛 Describe the bug

- Exact code worked fine on <= pytorch 2.2.1
- Breaks on pytorch > 2.2.1 (includes 2.3.0, 2.3.1)
- Deployment is with a docker container: nothing besides pytorch has changed (same CUDA version, etc). 
- `nvidia-smi` works fine.
- Without enviroment gating this works fine. 

Why is it necessary? As per [this thread](https://discuss.pytorch.org/t/is-it-expected-for-distributeddataparallel-to-use-more-memory-on-1-gpu-in-a-1gpu-1process-setup/77748/9) there is the possibility for memory leakage. This is more pronounced when using external libraries. Thus the cleanest way is to environment gate with `CUDA_VISIBLE_DEVICES`

Example: 
```bash
> python -c "import torch; print(torch.cuda.is_available()); print(torch.__version__)"
True
2.3.1+cu121

> nvidia-smi
nvidia-smi
Fri Jun  7 18:08:52 2024       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.182.03   Driver Version: 470.182.03   CUDA Version: 12.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA A100-SXM...  On   | 00000000:00:05.0 Off |                    0 |
| N/A   36C    P0    63W / 400W |      0MiB / 81251MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA A100-SXM...  On   | 00000000:00:06.0 Off |                    0 |
| N/A   34C    P0    67W / 400W |      0MiB / 81251MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
```

```bash
> ipython

In [1]: import torch
   ...: import os
   ...: torch.cuda.device_count()
   ...: os.environ["CUDA_VISIBLE_DEVICES"] = "0"
   ...: torch.empty(2, device='cuda')
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File /miniconda/lib/python3.10/site-packages/torch/cuda/__init__.py:306, in _lazy_init()
    305 try:
--> 306     queued_call()
    307 except Exception as e:

File /miniconda/lib/python3.10/site-packages/torch/cuda/__init__.py:174, in _check_capability()
    173 for d in range(device_count()):
--> 174     capability = get_device_capability(d)
    175     major = capability[0]

File /miniconda/lib/python3.10/site-packages/torch/cuda/__init__.py:430, in get_device_capability(device)
    418 r"""Get the cuda capability of a device.
    419 
    420 Args:
   (...)
    428     tuple(int, int): the major and minor cuda capability of the device
    429 """
--> 430 prop = get_device_properties(device)
    431 return prop.major, prop.minor

File /miniconda/lib/python3.10/site-packages/torch/cuda/__init__.py:448, in get_device_properties(device)
    447     raise AssertionError("Invalid device id")
--> 448 return _get_device_properties(device)

RuntimeError: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "../aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch. device=, num_gpus=

The above exception was the direct cause of the following exception:

DeferredCudaCallError                     Traceback (most recent call last)
Cell In[1], line 5
      3 torch.cuda.device_count()
      4 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
----> 5 torch.empty(2, device='cuda')

File /miniconda/lib/python3.10/site-packages/torch/cuda/__init__.py:312, in _lazy_init()
    307         except Exception as e:
    308             msg = (
    309                 f"CUDA call failed lazily at initialization with error: {str(e)}\n\n"
    310                 f"CUDA call was originally invoked at:\n\n{''.join(orig_traceback)}"
    311             )
--> 312             raise DeferredCudaCallError(msg) from e
    313 finally:
    314     delattr(_tls, "is_initializing")

DeferredCudaCallError: CUDA call failed lazily at initialization with error: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "../aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch. device=, num_gpus=

CUDA call was originally invoked at:

...
```

Without environment gating this works fine:
```bash
In [1]: import torch
   ...: import os
   ...: torch.cuda.device_count()
   ...: torch.empty(2, device='cuda')

Out[1]: tensor([0., 0.], device='cuda:0')
```

Related PR: https://github.com/pytorch/pytorch/issues/107300


### Versions

- Exact code worked fine on <= pytorch 2.2.1
- Breaks on pytorch > 2.2.1 (includes 2.3.0, 2.3.1)


cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @ptrblck