torch.amp

若API“是否支持”“是”“限制与说明”“-”,说明此API和原生API支持度保持一致。

API名称

NPU形式名称

是否支持

限制与说明

torch.autocast

torch.autocast

-

torch.cuda.amp.autocast

torch_npu.npu.amp.autocast

-

torch.cuda.amp.custom_fwd

torch_npu.npu.amp.custom_fwd

-

torch.cuda.amp.custom_bwd

torch_npu.npu.amp.custom_bwd

-

torch.cpu.amp.autocast

torch.cpu.amp.autocast

-

torch.cuda.amp.GradScaler

torch_npu.npu.amp.GradScaler

-