(self, args: tuple[Any, ...])
| 87 | return self.enabled |
| 88 | |
| 89 | def _resolve_device(self, args: tuple[Any, ...]) -> torch.device: |
| 90 | if self._device is not None: |
| 91 | return self._device |
| 92 | for arg in args: |
| 93 | if isinstance(arg, torch.Tensor) and arg.device.type == "cuda": |
| 94 | self._device = arg.device |
| 95 | return arg.device |
| 96 | self._device = torch.device("cuda") |
| 97 | return self._device |
| 98 | |
| 99 | def _ensure_stream(self, device: torch.device) -> torch.cuda.Stream: |
| 100 | if self._stream is None: |
no test coverage detected