在开启混合精度后，部分算子性能变差，如Reduce类算子

问题描述

开启混合精度后，Reduce类算子性能变差。

问题分析

由于block dim变小，性能降低。

处理方法

在AMP初始化之前添加代码：

amp.register_float_function(torch,'relu')      # relu为算子名称，请根据实际情况替换