amp用法没有对齐

[ccssu]

Summary

oneflow 和torch 对于 amp使用没有对齐

• 出现问题代码地点：https://github.com/Oneflow-Inc/one-yolov5/blob/fix_align_code/models/common.py#L650-L669

Code to reproduce bug

torch版代码

from torch.cuda import amp

autocast = False
with amp.autocast(autocast):
    print(autocast)
    print("is ok")

对应输出

(python3.8) fengwen@oneflow-25:~/one-yolov5$ /home/fengwen/miniconda3/envs/python3.8/bin/python /home/fengwen/one-yolov5/data/check_amp.py
False
is ok

oneflow版代码

from oneflow import amp

autocast = False
with amp.autocast(autocast):
    print(autocast)
    print("is ok")

对应输出

(python3.8) fengwen@oneflow-25:~/one-yolov5$ /home/fengwen/miniconda3/envs/python3.8/bin/python /home/fengwen/one-yolov5/data/check_amp.py
loaded library: /lib/x86_64-linux-gnu/libibverbs.so.1
Traceback (most recent call last):
  File "/home/fengwen/one-yolov5/data/check_amp.py", line 12, in <module>
    with amp.autocast(autocast):
  File "/home/fengwen/miniconda3/envs/python3.8/lib/python3.8/site-packages/oneflow/amp/autocast_mode.py", line 161, in __init__
    raise RuntimeError(
RuntimeError: User specified autocast device_type must be 'cuda' or 'cpu'

System Information

• What is your OneFlow installation (pip, source, dockerhub): pip install --pre oneflow -f https://staging.oneflow.info/branch/master/cu112
• OneFlow version (run python3 -m oneflow --doctor): '0.8.1.dev20221007+cu112'
• Python version: 3.8
• CUDA driver version:
  NVIDIA-SMI 515.65.01 Driver Version: 515.65.01 CUDA Version: 11.7
• 机器: oneflow25-root