diff --git a/muse_maskgit_pytorch/muse_maskgit_pytorch.py b/muse_maskgit_pytorch/muse_maskgit_pytorch.py
index af28e3f..5ba9597 100644
--- a/muse_maskgit_pytorch/muse_maskgit_pytorch.py
+++ b/muse_maskgit_pytorch/muse_maskgit_pytorch.py
@@ -510,7 +510,7 @@ def forward(
         # get loss
 
         ce_loss = self.transformer(
-            ids,
+            x,
             texts = texts,
             text_embeds = text_embeds,
             conditioning_token_ids = cond_token_ids,
diff --git a/setup.py b/setup.py
index 8c898c8..9b8228b 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'muse-maskgit-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.18',
+  version = '0.0.19',
   license='MIT',
   description = 'MUSE - Text-to-Image Generation via Masked Generative Transformers, in Pytorch',
   author = 'Phil Wang',