Repository: pytorch
Issue: NVML_SUCCESS == r INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":830, please report a bug to PyTorch. 
Labels: module: cuda,triaged,module: CUDACachingAllocator

### 🐛 Describe the bug

**Python sample code:**
```
from ultralytics.data.annotator import auto_annotate

auto_annotate(
data='/app/input/signal/images',
det_model='/app/runs/detect/train/weights/best.pt',
sam_model='mobile_sam.pt',
device="cuda",
output_dir='/app/input/signal/labels',
)
```

**error message:**

> RuntimeError                              Traceback (most recent call last)
> 
> Cell In[1], line 4
>       1 from ultralytics.data.annotator import auto_annotate
> ----> 4 auto_annotate(
>       5 data='/app/input/signal/images',
>       6 det_model='/app/runs/detect/train/weights/best.pt',
>       7 sam_model='mobile_sam.pt',
>       8 device="cuda",
>       9 output_dir='/app/input/signal/labels',
>      10 )
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/data/annotator.py:41, in auto_annotate(data, det_model, sam_model, device, output_dir)
>      39 if len(class_ids):
>      40     boxes = result.boxes.xyxy  # Boxes object for bbox outputs
> ---> 41     sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
>      42     segments = sam_results[0].masks.xyn  # noqa
>      44     with open(f'{str(Path(output_dir) / Path(result.path).stem)}.txt', 'w') as f:
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/model.py:35, in SAM.__call__(self, source, stream, bboxes, points, labels, **kwargs)
>      33 def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs):
>      34     """Calls the 'predict' function with given arguments to perform object detection."""
> ---> 35     return self.predict(source, stream, bboxes, points, labels, **kwargs)
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/model.py:31, in SAM.predict(self, source, stream, bboxes, points, labels, **kwargs)
>      29 kwargs.update(overrides)
>      30 prompts = dict(bboxes=bboxes, points=points, labels=labels)
> ---> 31 return super().predict(source, stream, prompts=prompts, **kwargs)
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/engine/model.py:242, in Model.predict(self, source, stream, predictor, **kwargs)
>     240 if prompts and hasattr(self.predictor, 'set_prompts'):  # for SAM-type models
>     241     self.predictor.set_prompts(prompts)
> --> 242 return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/engine/predictor.py:196, in BasePredictor.__call__(self, source, model, stream, *args, **kwargs)
>     194     return self.stream_inference(source, model, *args, **kwargs)
>     195 else:
> --> 196     return list(self.stream_inference(source, model, *args, **kwargs))
> 
> File /usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py:35, in _wrap_generator.<locals>.generator_context(*args, **kwargs)
>      32 try:
>      33     # Issuing `None` to a generator fires it up
>      34     with ctx_factory():
> ---> 35         response = gen.send(None)
>      37     while True:
>      38         try:
>      39             # Forward the response to our caller and get its next request
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/engine/predictor.py:259, in BasePredictor.stream_inference(self, source, model, *args, **kwargs)
>     257 # Inference
>     258 with profilers[1]:
> --> 259     preds = self.inference(im, *args, **kwargs)
>     261 # Postprocess
>     262 with profilers[2]:
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/predict.py:108, in Predictor.inference(self, im, bboxes, points, labels, masks, multimask_output, *args, **kwargs)
>     106 if all(i is None for i in [bboxes, points, masks]):
>     107     return self.generate(im, *args, **kwargs)
> --> 108 return self.prompt_inference(im, bboxes, points, labels, masks, multimask_output)
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/predict.py:170, in Predictor.prompt_inference(self, im, bboxes, points, labels, masks, multimask_output)
>     163 sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
>     164     points=points,
>     165     boxes=bboxes,
>     166     masks=masks,
>     167 )
>     169 # Predict masks
> --> 170 pred_masks, pred_scores = self.model.mask_decoder(
>     171     image_embeddings=features,
>     172     image_pe=self.model.prompt_encoder.get_dense_pe(),
>     173     sparse_prompt_embeddings=sparse_embeddings,
>     174     dense_prompt_embeddings=dense_embeddings,
>     175     multimask_output=multimask_output,
>     176 )
>     178 # (N, d, H, W) --> (N*d, H, W), (N, d) --> (N*d, )
>     179 # `d` could be 1 or 3 depends on `multimask_output`.
>     180 return pred_masks.flatten(0, 1), pred_scores.flatten(0, 1)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
>    1509     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
>    1510 else:
> -> 1511     return self._call_impl(*args, **kwargs)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
>    1515 # If we don't have any hooks, we want to skip the rest of the logic in
>    1516 # this function, and just call forward.
>    1517 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
>    1518         or _global_backward_pre_hooks or _global_backward_hooks
>    1519         or _global_forward_hooks or _global_forward_pre_hooks):
> -> 1520     return forward_call(*args, **kwargs)
>    1522 try:
>    1523     result = None
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/modules/decoders.py:79, in MaskDecoder.forward(self, image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings, multimask_output)
>      57 def forward(
>      58     self,
>      59     image_embeddings: torch.Tensor,
>    (...)
>      63     multimask_output: bool,
>      64 ) -> Tuple[torch.Tensor, torch.Tensor]:
>      65     """
>      66     Predict masks given image and prompt embeddings.
>      67 
>    (...)
>      77         torch.Tensor: batched predictions of mask quality
>      78     """
> ---> 79     masks, iou_pred = self.predict_masks(
>      80         image_embeddings=image_embeddings,
>      81         image_pe=image_pe,
>      82         sparse_prompt_embeddings=sparse_prompt_embeddings,
>      83         dense_prompt_embeddings=dense_prompt_embeddings,
>      84     )
>      86     # Select the correct mask or masks for output
>      87     mask_slice = slice(1, None) if multimask_output else slice(0, 1)
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/modules/decoders.py:118, in MaskDecoder.predict_masks(self, image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings)
>     115 b, c, h, w = src.shape
>     117 # Run the transformer
> --> 118 hs, src = self.transformer(src, pos_src, tokens)
>     119 iou_token_out = hs[:, 0, :]
>     120 mask_tokens_out = hs[:, 1:(1 + self.num_mask_tokens), :]
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
>    1509     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
>    1510 else:
> -> 1511     return self._call_impl(*args, **kwargs)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
>    1515 # If we don't have any hooks, we want to skip the rest of the logic in
>    1516 # this function, and just call forward.
>    1517 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
>    1518         or _global_backward_pre_hooks or _global_backward_hooks
>    1519         or _global_forward_hooks or _global_forward_pre_hooks):
> -> 1520     return forward_call(*args, **kwargs)
>    1522 try:
>    1523     result = None
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/modules/transformer.py:83, in TwoWayTransformer.forward(self, image_embedding, image_pe, point_embedding)
>      81 # Apply transformer blocks and final layernorm
>      82 for layer in self.layers:
> ---> 83     queries, keys = layer(
>      84         queries=queries,
>      85         keys=keys,
>      86         query_pe=point_embedding,
>      87         key_pe=image_pe,
>      88     )
>      90 # Apply the final attention layer from the points to the image
>      91 q = queries + point_embedding
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
>    1509     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
>    1510 else:
> -> 1511     return self._call_impl(*args, **kwargs)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
>    1515 # If we don't have any hooks, we want to skip the rest of the logic in
>    1516 # this function, and just call forward.
>    1517 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
>    1518         or _global_backward_pre_hooks or _global_backward_hooks
>    1519         or _global_forward_hooks or _global_forward_pre_hooks):
> -> 1520     return forward_call(*args, **kwargs)
>    1522 try:
>    1523     result = None
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/modules/transformer.py:153, in TwoWayAttentionBlock.forward(self, queries, keys, query_pe, key_pe)
>     151 q = queries + query_pe
>     152 k = keys + key_pe
> --> 153 attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
>     154 queries = queries + attn_out
>     155 queries = self.norm2(queries)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
>    1509     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
>    1510 else:
> -> 1511     return self._call_impl(*args, **kwargs)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
>    1515 # If we don't have any hooks, we want to skip the rest of the logic in
>    1516 # this function, and just call forward.
>    1517 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
>    1518         or _global_backward_pre_hooks or _global_backward_hooks
>    1519         or _global_forward_hooks or _global_forward_pre_hooks):
> -> 1520     return forward_call(*args, **kwargs)
>    1522 try:
>    1523     result = None
> 
> File /usr/local/lib/python3.10/dist-packages/ultralytics/models/sam/modules/transformer.py:212, in Attention.forward(self, q, k, v)
>     210 q = self.q_proj(q)
>     211 k = self.k_proj(k)
> --> 212 v = self.v_proj(v)
>     214 # Separate into heads
>     215 q = self._separate_heads(q, self.num_heads)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1511, in Module._wrapped_call_impl(self, *args, **kwargs)
>    1509     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
>    1510 else:
> -> 1511     return self._call_impl(*args, **kwargs)
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1520, in Module._call_impl(self, *args, **kwargs)
>    1515 # If we don't have any hooks, we want to skip the rest of the logic in
>    1516 # this function, and just call forward.
>    1517 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
>    1518         or _global_backward_pre_hooks or _global_backward_hooks
>    1519         or _global_forward_hooks or _global_forward_pre_hooks):
> -> 1520     return forward_call(*args, **kwargs)
>    1522 try:
>    1523     result = None
> 
> File /usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py:116, in Linear.forward(self, input)
>     115 def forward(self, input: Tensor) -> Tensor:
> --> 116     return F.linear(input, self.weight, self.bias)
> 
> RuntimeError: NVML_SUCCESS == r INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":830, please report a bug to PyTorch. 

### Versions

Collecting environment information...
PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-100-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A30
  MIG 1g.6gb      Device  0:

Nvidia driver version: 535.161.07
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             24
On-line CPU(s) list:                0-23
Vendor ID:                          AuthenticAMD
Model name:                         AMD EPYC 7313 16-Core Processor
CPU family:                         25
Model:                              1
Thread(s) per core:                 2
Core(s) per socket:                 6
Socket(s):                          2
Stepping:                           1
BogoMIPS:                           5988.78
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr rdpru arat umip vaes vpclmulqdq rdpid
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          384 KiB (12 instances)
L1i cache:                          384 KiB (12 instances)
L2 cache:                           6 MiB (12 instances)
L3 cache:                           64 MiB (2 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-11
NUMA node1 CPU(s):                  12-23
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET, no microcode
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.2.1
[pip3] torchvision==0.17.1
[pip3] triton==2.2.0
[conda] Could not collect

Click to add a cell.

cc @ptrblck