pytorch/torch/cuda
Matthew Hoffman e40d6ae0a7 Improve torch.cuda.amp type hints (#108630)
Fixes #108629

1. Add the following to their modules' `__all__` so that pyright considers them to be publicly exported:
* [`torch.autocast`](https://pytorch.org/docs/stable/amp.html#torch.autocast)
* [`torch.cuda.amp.GradScaler`](https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler)
* [`torch.cuda.amp.autocast`](https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast)
* [`torch.cuda.amp.custom_fwd`](https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.custom_fwd)
* [`torch.cuda.amp.custom_bwd`](https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.custom_bwd)
2. Add `overload`s for `torch.cuda.amp.GradScaler.scale` to differentiate when a `torch.Tensor` is returned vs. an `Iterable[torch.Tensor]` is returned based on the type of the `outputs` parameter.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/108630
Approved by: https://github.com/ezyang
2023-09-08 06:06:25 +00:00
..
amp Improve torch.cuda.amp type hints (#108630) 2023-09-08 06:06:25 +00:00
__init__.py Correctly format original traceback for delayed CUDA error (#107297) 2023-08-17 03:13:31 +00:00
_memory_viz.py MemoryViz.js: format, move style (#106482) 2023-08-03 00:42:13 +00:00
_sanitizer.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
_utils.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
comm.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
error.py
graphs.py Expose cudaStreamCaptureMode in CUDA Graphs, use local setting in inductor (#107407) 2023-08-25 01:44:26 +00:00
jiterator.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
memory.py [memory snapshots] _record_memory_history_legacy bug fix (#108260) 2023-08-30 22:33:04 +00:00
nccl.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
nvtx.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
profiler.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
random.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00
sparse.py
streams.py Apply UFMT to low traffic torch modules (#106249) 2023-07-29 23:37:30 +00:00