(beta)torch_npu.npu.set_mm_bmm_format_nd
接口原型
torch_npu.npu.set_mm_bmm_format_nd(bool)
功能描述
设置线性module里面的mm和bmm算子是否用ND格式。
支持的型号
Atlas 训练系列产品 Atlas A2 训练系列产品 Atlas A3 训练系列产品 Atlas 推理系列产品
调用示例
1 | >>>torch_npu.npu.set_mm_bmm_format_nd(True) |
父主题: torch_npu