fix: do not pass dtype to fourier embedding

Author: AshishKumar4

flaxdiff/models/simple_vit.py

@@ -127,8 +127,7 @@ def setup(self):
 
         # --- Conditioning ---
         self.time_embed = nn.Sequential([
-            # Fourier often uses float32
-            FourierEmbedding(features=self.emb_features, dtype=jnp.float32),
+            FourierEmbedding(features=self.emb_features),
             TimeProjection(features=self.emb_features,
                            dtype=self.dtype, precision=self.precision)
         ], name="time_embed")