NaN problem for fp16 UNet

Awesome work!
When I tried to run the code, I found that UNet sometimes output NaNs, and I found that this only happens when model dtype is fp16. To avoid this problem, simply use `.convert_to_fp32()` method of the UNet model.