Fix mixed precision for LearnableAffine

james77777778 · Mar 23, 2024 · e4a8eec · e4a8eec
1 parent 6ef977c
commit e4a8eec
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/kimm/layers/learnable_affine.py b/kimm/layers/learnable_affine.py
@@ -34,7 +34,9 @@ def build(self, input_shape):
         self.built = True
 
     def call(self, inputs, training=None, mask=None):
-        return ops.add(ops.multiply(inputs, self.scale), self.bias)
+        scale = ops.cast(self.scale, self.compute_dtype)
+        bias = ops.cast(self.bias, self.compute_dtype)
+        return ops.add(ops.multiply(inputs, scale), bias)
 
     def get_config(self):
         config = super().get_config()