训练的时候用的多GPU分布式训练,测试的时候没写分布式就报了这个错
测试的时候初始化模型:这个模型需要config参数
model = BertForSequenceClassification(config=bert_config)
分布式多GPU:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
加载训练好的模型:
output_model_file = './model/student_model.bin'
model.load_state_dict(torch.load(output_model_file))
放到GPU上测试
model.to(device)
model.eval()