fix: fixed a basic mistake

AshishKumar4 · Sep 8, 2024 · 6b9b4a4 · 6b9b4a4
1 parent 4ddb53c
commit 6b9b4a4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py
@@ -83,7 +83,7 @@ def __call__(self, x, temb, textcontext=None):
 
         # Patch embedding
         x = PatchEmbedding(patch_size=self.patch_size, embedding_dim=self.emb_features, 
-                           dtype=self.dtype, precision=self.precision)(x)
+                           dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init)(x)
         num_patches = x.shape[1]
 
         context_emb = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
@@ -98,7 +98,7 @@ def __call__(self, x, temb, textcontext=None):
         # print(f'Shape of x after time embedding: {x.shape}')
 
         # Add positional encoding
-        x = PositionalEncoding(max_len=x.shape[1], embedding_dim=self.emb_features, kernel_init=self.kernel_init)(x)
+        x = PositionalEncoding(max_len=x.shape[1], embedding_dim=self.emb_features)(x)
 
         # print(f'Shape of x after positional encoding: {x.shape}')
 

diff --git a/setup.py b/setup.py
@@ -11,7 +11,7 @@
 setup(
     name='flaxdiff',
     packages=find_packages(),
-    version='0.1.26',
+    version='0.1.27',
     description='A versatile and easy to understand Diffusion library',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',