diff --git a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py
index e570de5af..a0a9db008 100644
--- a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py
+++ b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py
@@ -5,7 +5,7 @@
 from torch.special import expm1
 
 from tqdm import tqdm
-from einops import rearrange, repeat
+from einops import rearrange, repeat, reduce
 from einops.layers.torch import Rearrange
 
 # helpers
@@ -268,7 +268,7 @@ def p_losses(self, x_start, times, noise = None):
         model_out = self.denoise_fn(x, log_snr)
 
         losses = self.loss_fn(model_out, noise, reduction = 'none')
-        losses = losses.mean(dim = tuple(range(1, losses.ndim)))
+        losses = reduce(losses, 'b ... -> b', 'mean')
 
         if self.p2_loss_weight_gamma >= 0:
             # following eq 8. in https://arxiv.org/abs/2204.00227
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index e9115c8fb..1c9fbc93e 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -541,7 +541,7 @@ def forward(self, img, *args, **kwargs):
 # dataset classes
 
 class Dataset(data.Dataset):
-    def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png']):
+    def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png'], augment_horizontal_flip = False):
         super().__init__()
         self.folder = folder
         self.image_size = image_size
@@ -549,7 +549,7 @@ def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png']):
 
         self.transform = transforms.Compose([
             transforms.Resize(image_size),
-            transforms.RandomHorizontalFlip(),
+            transforms.RandomHorizontalFlip() if augment_horizontal_flip else nn.Identity(),
             transforms.CenterCrop(image_size),
             transforms.ToTensor()
         ])
@@ -580,7 +580,8 @@ def __init__(
         step_start_ema = 2000,
         update_ema_every = 10,
         save_and_sample_every = 1000,
-        results_folder = './results'
+        results_folder = './results',
+        augment_horizontal_flip = True
     ):
         super().__init__()
         self.model = diffusion_model
@@ -596,7 +597,7 @@ def __init__(
         self.gradient_accumulate_every = gradient_accumulate_every
         self.train_num_steps = train_num_steps
 
-        self.ds = Dataset(folder, image_size)
+        self.ds = Dataset(folder, image_size, augment_horizontal_flip = augment_horizontal_flip)
         self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle=True, pin_memory=True))
         self.opt = Adam(diffusion_model.parameters(), lr=train_lr)
 
diff --git a/setup.py b/setup.py
index 3c62cbbae..955cc9de8 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.18.0',
+  version = '0.18.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',