Update model.py

model.py

@@ -89,6 +89,7 @@ class PatchEmbedding(nnx.Module):
             padding=config.padding,
             use_bias=config.use_bias,
             rngs=rngs,
+            dtype=jnp.bfloat16,
         )
 
     def __call__(self, x):
@@ -103,10 +104,10 @@ class TimeEmbedding(nnx.Module):
         self.freq_dim = config.time_freq_dim
         self.max_period = config.time_max_period
         self.fc1 = nnx.Linear(
-            self.freq_dim, config.hidden_dim, use_bias=config.use_bias, rngs=rngs
+            self.freq_dim, config.hidden_dim, use_bias=config.use_bias, rngs=rngs, dtype=jnp.bfloat16
         )
         self.fc2 = nnx.Linear(
-            config.hidden_dim, config.hidden_dim, use_bias=config.use_bias, rngs=rngs
+            config.hidden_dim, config.hidden_dim, use_bias=config.use_bias, rngs=rngs, dtype=jnp.bfloat16
         )
 
     @staticmethod
@@ -140,12 +141,14 @@ class MLP(nnx.Module):
             config.hidden_dim * config.mlp_ratio,
             use_bias=config.use_bias,
             rngs=rngs,
+            dtype=jnp.bfloat16,
         )
         self.fc2 = nnx.Linear(
             config.hidden_dim * config.mlp_ratio,
             config.hidden_dim,
             use_bias=config.use_bias,
             rngs=rngs,
+            dtype=jnp.bfloat16,
         )
 
     def __call__(self, x):
@@ -165,6 +168,7 @@ class SelfAttention(nnx.Module):
             3 * config.hidden_dim,
             use_bias=config.use_bias,
             rngs=rngs,
+            dtype=jnp.bfloat16,
         )
         self.heads = config.num_heads
         self.head_dim = config.hidden_dim // config.num_heads
@@ -209,6 +213,7 @@ class TransformerBlock(nnx.Module):
                 6 * config.hidden_dim,
                 use_bias=config.use_bias,
                 rngs=rngs,
+                dtype=jnp.bfloat16,
             ),
         )
 
@@ -241,6 +246,7 @@ class FinalLayer(nnx.Module):
             padding=config.padding,
             use_bias=config.use_bias,
             rngs=rngs,
+            dtype=jnp.bfloat16,
         )
         self.adalm_modulation = nnx.Sequential(
             nnx.silu,
@@ -249,6 +255,7 @@ class FinalLayer(nnx.Module):
                 2 * config.hidden_dim,
                 use_bias=config.use_bias,
                 rngs=rngs,
+                dtype=jnp.bfloat16,
             ),
         )