From d97bc0278c7254e898ac27a237f1bb743807f240 Mon Sep 17 00:00:00 2001
From: Arya Aftab <arya.aftab@gmail.com>
Date: Wed, 29 Jun 2022 11:42:54 +0430
Subject: [PATCH] Solve issue #26

---
 denoising_diffusion_pytorch/denoising_diffusion_pytorch.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index e5fce00..ac83a3c 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -15,12 +15,17 @@ from torch.optim import Adam
 from torchvision import transforms, utils
 from PIL import Image
 
-from tqdm import tqdm
 from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
 
 from ema_pytorch import EMA
 
+import sys
+if 'ipykernel' in sys.modules:
+    from tqdm.notebook import tqdm
+else:
+    from tqdm import tqdm
+
 # helpers functions
 
 def exists(x):