Michael Carilli 3b040c478a Make custom_fwd a no-op when not executed under autocast (#36171) ... Summary: Currently, a custom autograd function written with ``` torch.cuda.amp.custom_fwd(cast_inputs=dtype) def forward(ctx, *args): ... ``` casts incoming floating-point CUDA tensors to `dtype` unconditionally, regardless of whether the function executes in an autocast-enabled region. I think I had the wrong idea there. Autocast-disabled regions should give the user control of input types. Also, `custom_fwd(cast_inputs=dtype)`-decorated functions' behavior should align with native fp32list/fp16list functions. C++-side casting wrappers have no effect when autocast is disabled, and `custom_fwd`'s casting should behave the same way. The present PR changes `custom_fwd` so it only casts in autocast-enabled regions (also updates custom_fwd to ignore fp64 inputs, like the C++ wrappers). Pull Request resolved: https://github.com/pytorch/pytorch/pull/36171 Differential Revision: D22179511 Pulled By: ngimel fbshipit-source-id: 5a93d070179a43206066bce19da0a5a19ecaabbd		2020-06-23 10:23:02 -07:00
..
__init__.py	[RELAND2] Eager autocasting, out-of-place ops only (with MSVC 2017 fix) (#35102)	2020-03-24 09:08:04 -07:00
autocast_mode.py	Make custom_fwd a no-op when not executed under autocast (#36171)	2020-06-23 10:23:02 -07:00
grad_scaler.py	Allow GradScaler to be pickled (#38296)	2020-05-14 09:14:28 -07:00