Add param_dtype to AddPositionEmbs.

PiperOrigin-RevId: 615347924
google-research · Mar 13, 2024 · 1ff0176 · 1ff0176
1 parent 3d8e814
commit 1ff0176
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/vit_jax/models_vit.py b/vit_jax/models_vit.py
@@ -42,6 +42,7 @@ class AddPositionEmbs(nn.Module):
   """
 
   posemb_init: Callable[[PRNGKey, Shape, Dtype], Array]
+  param_dtype: Dtype = jnp.float32
 
   @nn.compact
   def __call__(self, inputs):
@@ -57,7 +58,8 @@ def __call__(self, inputs):
     assert inputs.ndim == 3, ('Number of dimensions should be 3,'
                               ' but it is: %d' % inputs.ndim)
     pos_emb_shape = (1, inputs.shape[1], inputs.shape[2])
-    pe = self.param('pos_embedding', self.posemb_init, pos_emb_shape)
+    pe = self.param(
+        'pos_embedding', self.posemb_init, pos_emb_shape, self.param_dtype)
     return inputs + pe