pytorch中数据和模型都要部署在cuda上面

热爱生活的五柒2024-11-15 6:00

注意一定要两个都部署在模型训练之前，缺少一个都会报错在两个设备上训练，一个在cpu，一个在cuda

部署数据到cuda上

复制代码

point_features = torch.tensor(point_features, dtype=torch.float32).to('cuda')  # 如果你有支持的 GPU,shape = (499，3159,3)
coord_time = torch.tensor(coord_time, dtype=torch.float32).to('cuda')  # 如果你有支持的 GPU  shape = (3159,3)
target = torch.tensor(target, dtype=torch.float32).to('cuda')  # 如果你有支持的 GPU  shape = (499，3159,3)

部署模型到cuda上

model = model.to('cuda')