LayerConvNeXtGamma: fix for mixed precision mode

This commit is contained in:
Starbeamrainbowlabs 2023-02-02 16:22:08 +00:00
parent a630db2c49
commit 1a8f10339a
Signed by: sbrl
GPG key ID: 1BE5172E637709C2

View file

@ -7,7 +7,7 @@ class LayerConvNeXtGamma(tf.keras.layers.Layer):
super(LayerConvNeXtGamma, self).__init__(name=name)
self.dim = dim
self.const = const_val * tf.ones((self.dim))
self.const = const_val * tf.ones((self.dim), dtype=tf.float32 if tf.mixed_precision.global_policy().name == "float32" else tf.float16)
def call(self, inputs, **kwargs):
return tf.multiply(inputs, self.const)