pytorch框架中如何在推理过程中使用多gpu

我正试图从基于pytorch框架的unet3D进行模型预测。我正在使用多GPU

import torch
import os
import torch.nn as nn
os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES']='0,1,2'

model = unet3d()
model = nn.DataParallel(model)
model = model.to('cuda')

result = model.forward(torch.tensor(input).to('cuda').float())

但该模型仍然只使用1个GPU（第一个），我得到内存错误

CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 11.00 GiB total capacity; 8.43 GiB already allocated; 52.21 MiB free; 5.17 MiB cached) 

---

在推理阶段，我应该如何使用多个GPU？我上面的脚本有什么错误

DataParallel处理将数据发送到gpu的过程

import torch
import os
import torch.nn as nn
os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES']='0,1,2'

model = unet3d()
model = nn.DataParallel(model.cuda())

result = model.forward(torch.tensor(input).float())

如果这不起作用，请提供有关

输入的更多详细信息

[编辑]：

试试这个：

with torch.no_grad():
    result = model(torch.tensor(input).float())

您打算如何在批处理大小为1的情况下进行多GPU推断？Pytork应该如何跨GPU分割数据？这是不可能“开箱即用”的。（它应该跨GPU分割模型吗？它应该将图像分割成两半还是四分之一？）。在keras中也会发生同样的情况，方法是将权重和层分布在不同的GPU上。请编辑您的问题，以提供一个keras片段来演示此功能。