How to use float32 for a particular nn.Module?

[turian]

I want to run most of my training in bfloat16.

However, I have one nn.Module that has a forward method receiving an int (its seed) and then it does a bunch of ops that are not bfloat16 supported. I'd like for just this module to work in float32, and then cast it to bfloat16 at the end.

I don't think this will slow things down. But how do I do it elegantly?