Update tinydit.py

tinydit.py

@@ -90,7 +90,7 @@ class TinyDit(nn.Module):
           nn.GELU(),
           nn.Linear(dim, dim)
       )
-      self.pos_embed = nn.Parameter(torch.randn(1, 16, dim))
+      self.pos_embed = nn.Parameter(torch.randn(1, latent_channels, dim))
       self.blocks = nn.ModuleList([DiTBlock(dim=dim, heads=8, mlp_dim=dim*4) for _ in range(depth)])
       self.unpatchify = Unpatchify(patch_size=patch_size, dim=dim, out_channels=latent_channels)