fix

Created using spr 1.3.4

Author: boomanaiden154

gematria/model/python/model_blocks.py

@@ -78,3 +78,13 @@ def call(self, layer_inputs: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
       residual_part = self._linear_transformation(residual_part)
 
     return tf.math.add(output_part, residual_part, name=self.name)
+
+
+class CastLayer(tf_keras.layers.Layer):
+
+  def __init__(self, dtype, **kwargs):
+    super().__init__(**kwargs)
+    self._dtype = dtype
+
+  def call(self, input_tensor):
+    return tf.cast(input_tensor, dtype=self._dtype)