Error with custom normalization layer

Hi, thank you for your great work.
I have a question regarding custom normalization layer.

Since all of the modules which is not
'isinstance(module, torch.nn.modules.batchnorm._BatchNorm) and module.affine is True'
will be converted to fp16, custom normalization layer such as batch-instance-normalization(https://github.com/hyeonseob-nam/Batch-Instance-Normalization/blob/master/models/batchinstancenorm.py) throws an 'RuntimeError: expected scalar type Half but found Float'.

I have tried converting input into float only in custom norm layer to resolve the error above. However, it resulted in more GPU memory usage than usual fp32 training.
What is the best way to unconvert those module after amp.initialize?

Also, as I debugged batchnorm layer, I have found that input tensor to batchnorm is half type on opt_level 'O1'. Does 'keep_batchnorm_fp32' means only keeping weight of batchnorm in fp32?

Thank you for your help.

Start at the amp.initialize entry point and inspect the module-conversion logic described in the issue, then reproduce the error with the linked batchinstancenorm.py custom layer. Check the O1 and keep_batchnorm_fp32 paths, including the observed BatchNorm input type. Done means the custom normalization behavior and BatchNorm precision semantics are clarified or corrected with a regression test.