casting the unet to low precision when training controlnet

[eyalgutflaish]

I see that the unet is casted to low precision on examples/controlnet/train_controlnet_sdxl.py.
As I understand when using automatic mixed precision by accelerate, all layers that propagate gradients should be left on full precision and the accelerate module should decide automatically on the precision of layers.

So why do we cast the unet to low precision ?

[tolgacangoz]

According to accelerate's documentation, "As a rule of thumb, for stable training with automatic mixed precision, all the trainable parameters have to be in torch.float32.".

[eyalgutflaish]

But what about non-trainable parameters that propagates gradients. They might have layers that operate better in full precision. Shouldn't automatic mixed precision of accelerate "decide" on them ?