tensorplay.amp.autocast_mode.custom_fwd

tensorplay.amp.autocast_mode.custom_fwd(fwd=None, *, device_type: str, cast_inputs=None)[source]

Create a helper decorator for forward methods of custom autograd functions.

Autograd functions are subclasses of tensorplay.autograd.Function.

Parameters:

• device_type (str) – Device type to use. ‘cuda’, ‘cpu’. The type is the same as the type attribute of a tensorplay.device. Thus, you may obtain the device type of a tensor using Tensor.device.type.

• cast_inputs (tensorplay.dtype or None, optional, default=None) – If not None, when forward runs in an autocast-enabled region, casts incoming floating-point Tensors to the target dtype (non-floating-point Tensors are not affected), then executes forward with autocast disabled. If None, forward’s internal ops execute with the current autocast state.

Note

If the decorated forward is called outside an autocast-enabled region, custom_fwd is a no-op and cast_inputs has no effect.