pytorch将模型加载到GPU上
# ----------- 判断模型是在CPU还是GPU上 ---------------------- model = nn.LSTM(input_size=10, hidden_size=4, num_layers=1, batch_first=True)print(next(model.parameters()).device) # 输出:cpu model = model.cuda()print(next(model.parameters()).device) # 输出:cuda:.
原创
2022-02-14 21:22:04 ·
6073 阅读 ·
0 评论