(gpu_to_use: int | None, torch_kwargs: dict)
| 1979 | |
| 1980 | |
| 1981 | def _update_device(gpu_to_use: int | None, torch_kwargs: dict) -> None: |
| 1982 | if "device" not in torch_kwargs and gpu_to_use is not None: |
| 1983 | device = _gpu_to_use_to_device(gpu_to_use, device=None) |
| 1984 | if device is not None: |
| 1985 | torch_kwargs["device"] = device |
| 1986 | |
| 1987 | |
| 1988 | def _gpu_to_use_to_device(gpu_to_use: int | None, device: str | None) -> str | None: |
no test coverage detected