Create UNet that uses float16 weight

[tiwwexx]

Hello,

I'm trying to train a 3D unet using monai's UNet model. I've been getting a bunch of out of memory issues and thought that using 16bit instead of 32bit inputs and weights might help. Anyone have an Idea on how to get the UNet to be float16,
something like UNet(... , astype = 'Halftensor' )

Any help is much appreciated!

[Nic-Ma]

Hi @tiwwexx ,

I think you can check the AMP tutorial:
https://github.com/Project-MONAI/tutorials/blob/master/acceleration/automatic_mixed_precision.ipynb

Thanks.