关于train.py

[yoho131]

作者你好，我在尝试训练我自己的数据集的时候出现了一些小问题。
当我尝试使用train.sh脚本的时候设备会重启，当我单独运行train.py的时候会报错如下：

能否解惑，万分感谢

[zqhang]

这个应该是text_prob和label的size不匹配，你打印这两个变量的shape看看

[LJK12581]

提问者您好,请问这个问题解决了嘛？

[yoho131]

我将train.py的83行
image_loss = F.cross_entropy(text_probs.squeeze(), label.long().cuda())
修改为
image_loss = F.cross_entropy(text_prob, label.long().cuda())
之后可以成功运行，但是我不知道这是否会对训练结果带来不利影响