Windows bug ModuleNotFoundError: No module named 'torch._six'

[MGzhou]

System Info

包版本如下

torch                       2.0.1
flagai                      1.8.4

在使用 https://model.baai.ac.cn/model-detail/100007 的模型时弹出这个bug。

经过查询，发现原因是因为 torch 升级到2.0之后，这个api不能用了。

经过查看源码，"flagai\mpu\grads.py"

import torch
import sys
if torch.__version__ >= "2" and sys.platform != "win32":
    from torch import inf
else:
    from torch._six import inf

发现代码已经针对这个api已经改进，但排除了win32 ，即windows系统。

我去掉 and sys.platform != "win32" 后代码正确运行

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as T5/AltCLIP, ...)
• My own task or dataset (give details below)

Reproduction

import torch
from flagai.auto_model.auto_loader import AutoLoader
from flagai.model.predictor.predictor import Predictor

device = torch.device("cpu")
maxlen = 256

###
auto_loader = AutoLoader(
    task_name="title-generation",
    model_name="RoBERTa-base-ch",   # Load the checkpoints from Modelhub(model.baai.ac.cn/models),
    device=device,
)

model = auto_loader.get_model()
model.to(device)
tokenizer = auto_loader.get_tokenizer()

predictor = Predictor(model, tokenizer)

test_data = [
    "本文总结了十个可穿戴产品的设计原则而这些原则同样也是笔者认为是这个行业最吸引人的地方1为人们解决重复性问题2从人开始而不是从机器开始3要引起注意但不要刻意4提升用户能力而不是取代人",
    "2007年乔布斯向人们展示iPhone并宣称它将会改变世界还有人认为他在夸大其词然而在8年后以iPhone为代表的触屏智能手机已经席卷全球各个角落未来智能手机将会成为真正的个人电脑为人类发展做出更大的贡献",
    "雅虎发布2014年第四季度财报并推出了免税方式剥离其持有的阿里巴巴集团15％股权的计划打算将这一价值约400亿美元的宝贵投资分配给股东截止发稿前雅虎股价上涨了大约7％至5145美元"
]

for text in test_data:
    print(
        predictor.predict_generate_beamsearch(text,
                                              out_max_length=50,
                                              beam_size=3))

Expected behavior

如果这个bug是真实的，希望修复这个bug。