import torch
# 使用cpu训练
# device = torch.device('cpu')
# 使用GPU训练,当电脑中有多张显卡时，使用第一张显卡
device = torch.device('cuda:0')

# 模型转到device
model=model.to(device)

# 训练数据转到device
input_batch = input_batch.to(device)
target_batch = target_batch.to(device)

# 训练完成进行模型存储
# 模型转到cpu
model.to('cpu')
# 模型路径
modelFilePath_str = 'model\\testModel.model'
# 存储模型
# torch.save(model,modelFilePath_str)
# 加载模型
testModel=torch.load(modelFilePath_str)