diff --git a/autoregressive_diffusion_pytorch/autoregressive_diffusion.py b/autoregressive_diffusion_pytorch/autoregressive_diffusion.py
index f66b218..b4d1b7c 100644
--- a/autoregressive_diffusion_pytorch/autoregressive_diffusion.py
+++ b/autoregressive_diffusion_pytorch/autoregressive_diffusion.py
@@ -511,6 +511,12 @@ def forward(
 
 # image wrapper
 
+def normalize_to_neg_one_to_one(img):
+    return img * 2 - 1
+
+def unnormalize_to_zero_to_one(t):
+    return (t + 1) * 0.5
+
 class ImageAutoregressiveDiffusion(Module):
     def __init__(
         self,
@@ -541,8 +547,10 @@ def __init__(
 
     def sample(self, batch_size = 1):
         tokens = self.model.sample(batch_size = batch_size)
-        return self.to_image(tokens)
+        images = self.to_image(tokens)
+        return unnormalize_to_zero_to_one(images)
 
     def forward(self, images):
+        images = normalize_to_neg_one_to_one(images)
         tokens = self.to_tokens(images)
         return self.model(tokens)
diff --git a/pyproject.toml b/pyproject.toml
index 96d50a9..40081ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "autoregressive-diffusion-pytorch"
-version = "0.0.5"
+version = "0.0.6"
 description = "Autoregressive Diffusion - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }