How to use slide_window_inference with enough gpu memory but limited cpu memory

[uto-lt]

Hello, MONAI team!
Thanks for the convience your software provides so far!
I tried to develop an algorithm for segmentation of abdomen organs to a big dataset under docker. But it comes to an memory error when I infer a test volume with slide_window_inference, a 512x512x491 CT volume in nifti format.
I have enough gpu memory(48G, A6000), but I just have 28G cpu memory, which is limited by the competition organizer. So, I want to make full use of gpu, while use cpu memory as little as possible.
I have put both the data and model into cuda by the following code:

model = model.to(device)
images = test_data["image"].to(device)
roi_size = (96, 96, 96)
sw_batch_size = 4
test_data['pred'] = sliding_window_inference(images, roi_size, sw_batch_size, model, overlap=0.5)

But it seems that the peak cpu memory is more than 28GB, while the gpu memory is around 10GB.
What should I do next? Looking forward for your help!

[wyli]

perhaps you can consider a lower precision on CPU, for example images = images.to(torch.float16)

[uto-lt]

Do you meanimages = test_data["image"].to(torch.float16).to(device)?
But it doesn't work :(