Use correct dtype for transformer trainer init

neurobayes/flax_nets/deterministic_nn.py

@@ -49,14 +49,21 @@ def __init__(self,
 
         input_shape = (input_shape,) if isinstance(input_shape, int) else input_shape
         self.model = architecture
+
+        is_transformer = any(base.__name__.lower().find('transformer') >= 0 
+                        for base in architecture.__mro__)
+        input_dtype = jnp.int32 if is_transformer else jnp.float32
 
         if loss not in ['homoskedastic', 'heteroskedastic', 'classification']:
             raise ValueError("Select between 'homoskedastic', 'heteroskedastic', or 'classification' loss")
         self.loss = loss
 
         # Initialize model
         key = jax.random.PRNGKey(0)
-        params = self.model.init(key, jnp.ones((1, *input_shape)))['params']
+        params = self.model.init(
+            key,
+            jnp.ones((1, *input_shape), dtype=input_dtype)
+        )['params']
 
         # Default SWA configuration with all required parameters
         self.default_swa_config = {