I use the Unet3D model from pytorch-unet3d repo. I am training/predict on the machine with multiple gpus, but use only one GPU device. For training, I use the code provided by the repository but use my own dataloaders. For the prediction I use the following code
device = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")
config['device'] = device
model = get_model(config)
utils.load_checkpoint(PATH, model)

with torch.no_grad():
    model.to(device)
    model.eval()
    arr = arr.to(device)
    pred = model(arr)
    pred = pred.data.cpu().numpy()
    print(np.unique(pred),pred.shape)

result of print(model)
....
....
)
  (final_conv): Conv3d(32, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))
  (final_activation): Sigmoid()
)

Versions
Pytorch=1.2, Cuda compilation tools, release 10.0, V10.0.130
pytorch was installed via  pip install torch==1.2.0 torchvision==0.4.0
I don’t know how to check cudnn version, could you please help me with that?