more fixes

AshishKumar4 · AshishKumar4 · commit 4938f552ec03 · 2025-05-05T14:49:34.000Z
diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py
@@ -295,8 +295,8 @@ def setup(self):
             )
 
         self.time_embed = nn.Sequential([
-            FourierEmbedding(features=self.emb_features, dtype=jnp.float32),
-            TimeProjection(features=self.emb_features * self.mlp_ratio, dtype=self.dtype, precision=self.precision),
+            FourierEmbedding(features=self.emb_features),
+            TimeProjection(features=self.emb_features * self.mlp_ratio),
             nn.Dense(features=self.emb_features, dtype=self.dtype, precision=self.precision)
         ], name="time_embed")
 

Original file line number	Diff line number	Diff line change
`@@ -295,8 +295,8 @@ def setup(self):`
`295`	`295`	`)`
`296`	`296`
`297`	`297`	`self.time_embed = nn.Sequential([`
`298`		`- FourierEmbedding(features=self.emb_features, dtype=jnp.float32),`
`299`		`- TimeProjection(features=self.emb_features * self.mlp_ratio, dtype=self.dtype, precision=self.precision),`
	`298`	`+ FourierEmbedding(features=self.emb_features),`
	`299`	`+ TimeProjection(features=self.emb_features * self.mlp_ratio),`
`300`	`300`	`nn.Dense(features=self.emb_features, dtype=self.dtype, precision=self.precision)`
`301`	`301`	`], name="time_embed")`
`302`	`302`