From 1a8f10339afd1eb826cd5254e43a5b702a2db72d Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Thu, 2 Feb 2023 16:22:08 +0000 Subject: [PATCH] LayerConvNeXtGamma: fix for mixed precision mode --- aimodel/src/lib/ai/components/LayerConvNeXtGamma.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aimodel/src/lib/ai/components/LayerConvNeXtGamma.py b/aimodel/src/lib/ai/components/LayerConvNeXtGamma.py index 9f29564..1189de5 100644 --- a/aimodel/src/lib/ai/components/LayerConvNeXtGamma.py +++ b/aimodel/src/lib/ai/components/LayerConvNeXtGamma.py @@ -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)