frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Replace device check of TORCH_INTERNAL_ASSERT with TORCH_CHECK (#155318)

Browse Source 
Fixes #136849

## Test Result

```python
>>> import torch
>>> device = torch.cuda.device_count() + 1
>>> torch.cuda.current_stream(device) #  INTERNAL ASSERT FAILED
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/zong/code/pytorch/torch/cuda/__init__.py", line 1083, in current_stream
    streamdata = torch._C._cuda_getCurrentStream(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Device index value 3 is out of index range [0, 2)

>>> torch.cuda.default_stream(device) #  INTERNAL ASSERT FAILED
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/zong/code/pytorch/torch/cuda/__init__.py", line 1101, in default_stream
    streamdata = torch._C._cuda_getDefaultStream(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Device index value 3 is out of index range [0, 2)

>>> torch.cuda.set_per_process_memory_fraction(0.5, device)  #  INTERNAL ASSERT FAILED
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/zong/code/pytorch/torch/cuda/memory.py", line 193, in set_per_process_memory_fraction
    torch._C._cuda_setMemoryFraction(fraction, device)
RuntimeError: Allocator not initialized for device : did you call init?

```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/155318
Approved by: https://github.com/albanD
This commit is contained in:
zeshengzong
2025-06-13 01:20:19 +00:00
committed by PyTorch MergeBot
parent d7e657da35
commit 905b194a2e
3 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
test/test_cuda.py
View File

@ -1061,7 +1061,7 @@ print(t.is_pinned())
torch.accelerator.set_stream(s2)
self.assertEqual(torch.accelerator.current_stream().stream_id, s2.stream_id)
with self.assertRaisesRegex(
RuntimeError, "device_index >= 0 && device_index < num_gpus"
RuntimeError, "Device index value .* is out of index range"
):
torch.accelerator.current_stream(torch.accelerator.device_count())