tensorplay.amp.autocast_mode.custom_bwd

tensorplay.amp.autocast_mode.custom_bwd(bwd=None, *, device_type: str)[source]

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

Autograd functions are subclasses of tensorplay.autograd.Function. Ensures that backward executes with the same autocast state as forward.

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.