pytorch/torch/utils
shibo 6aeb85add8 add checkpoint support for custom device (#99626)
Fixes #ISSUE_NUMBER
1、add checkpoint support for custom device
2、add a device argument, I want to add a device="cuda" parameter to the func `forward` of `CheckpointFunction`, and I can specify the device type when using it, but the func `apply` of `torch.autograd.Function` does not support `kwargs`, so I added a variable named `_device`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/99626
Approved by: https://github.com/soulitzer
2023-05-04 00:23:42 +00:00
..
_sympy [export] Constraints API (#98433) 2023-04-13 21:20:10 +00:00
backcompat [BE]: pyupgrade Python to 3.8 - imports and object inheritance only (#94308) 2023-02-07 21:10:56 +00:00
benchmark torch.compile benchmark utility (#97699) 2023-04-12 03:02:06 +00:00
bottleneck Fix use-dict-literal lint (#83718) 2022-08-24 00:26:46 +00:00
data [DataPipe] Change DataPipe display name in profiler (#100042) 2023-05-03 21:36:13 +00:00
hipify [BE] Fix SIM109 compare-with-tuple (#100337) 2023-04-30 19:51:32 +00:00
jit Create __init__.py (#78629) 2022-06-03 18:14:21 +00:00
model_dump [BE]: Merge startswith calls - rule PIE810 (#96754) 2023-03-14 22:05:20 +00:00
tensorboard fix DeprecationWarning (#95545) 2023-02-28 03:43:57 +00:00
__init__.py Automatically generate attributes and methods for custom backends. (#98066) 2023-04-13 22:04:05 +00:00
_content_store.py Make hash_storage work with size 0/1 storage (#100467) 2023-05-03 03:10:19 +00:00
_contextlib.py Reland "Add torch.utils.device_mode" (#91796) 2023-01-09 20:57:12 +00:00
_cpp_extension_versioner.py [BE]: pyupgrade Python to 3.8 - imports and object inheritance only (#94308) 2023-02-07 21:10:56 +00:00
_crash_handler.py Use CMake for breakpad (#63186) 2021-08-19 10:42:01 -07:00
_cuda_trace.py Don't use f-strings in logging calls (1/X) (#98591) 2023-04-07 15:52:50 +00:00
_device.py Guard on default device (#99551) 2023-04-20 17:02:59 +00:00
_foreach_utils.py [BE] Enable C419 rule for any all shortcircuiting (#99890) 2023-04-25 15:02:13 +00:00
_freeze.py [BE] Enable C419 rule for any all shortcircuiting (#99890) 2023-04-25 15:02:13 +00:00
_mode_utils.py [Modes] remove enable and rewrite mode stack (squashed) (#84774) 2022-09-27 01:04:35 +00:00
_python_dispatch.py fix per-dispatchkey-mode caching bug (#98030) 2023-04-25 21:58:14 +00:00
_pytree.py Improve code to make it more pythonic. (#99720) 2023-04-23 16:42:13 +00:00
_stats.py Fast path binary ops in fake tensor (#94047) 2023-02-07 18:34:24 +00:00
_traceback.py Add utility report_compile_source_on_error (#91069) 2023-01-11 22:54:46 +00:00
_zip.py [BE] Prefer dash over underscore in command-line options (#94505) 2023-02-09 20:16:49 +00:00
backend_registration.py Fix a minor bug about method generation. (#99704) 2023-04-24 22:18:18 +00:00
bundled_inputs.py Deprecate TypedStorage, its derived classes, and all of their public methods (#85303) 2022-11-08 18:11:01 +00:00
checkpoint.py add checkpoint support for custom device (#99626) 2023-05-04 00:23:42 +00:00
collect_env.py [BE] Fix collect_env for python-path-with-space (#98415) 2023-04-06 01:09:23 +00:00
cpp_backtrace.py Expose cpp_backtrace to python binding (#84896) 2022-09-27 14:59:08 +00:00
cpp_extension.py Use correct standard when compiling NVCC on Windows (#100031) 2023-05-01 16:28:23 +00:00
dlpack.py Add DLPack support for XPU backend by mapping to kDLOneAPI in DLPack … (#94968) 2023-03-30 04:32:15 +00:00
file_baton.py Fix typos under torch/utils directory (#97516) 2023-03-24 16:53:39 +00:00
flop_counter.py tweak heuristic for sdpa selection based off of *data* (and a decision tree) (#99644) 2023-04-21 23:28:44 +00:00
hooks.py Fix module backward pre-hooks to actually update gradient (#97983) 2023-03-30 20:33:44 +00:00
mkldnn.py [BE] [2/3] Rewrite super() calls in functorch and torch (#94588) 2023-02-10 21:16:33 +00:00
mobile_optimizer.py [BE] Enable C419 rule for any all shortcircuiting (#99890) 2023-04-25 15:02:13 +00:00
model_zoo.py
show_pickle.py [BE]: pyupgrade Python to 3.8 - imports and object inheritance only (#94308) 2023-02-07 21:10:56 +00:00
throughput_benchmark.py [BE]: pyupgrade Python to 3.8 - imports and object inheritance only (#94308) 2023-02-07 21:10:56 +00:00
weak.py cudagraph trees (#89146) 2023-03-17 02:47:03 +00:00