Generated samples contain only noise while training unconditional Imagen

import torch
from imagen_pytorch import Unet, Imagen, ImagenTrainer
from tqdm import tqdm
import torch.optim as optim
import torchvision
from torchvision import transforms as T

class CifarDataset(Dataset):
    def __init__(self):
        image_size = 32
        super().__init__()
        self.transform = T.Compose([
            T.Resize(image_size),
            T.RandomHorizontalFlip(),
            T.CenterCrop(image_size),
            T.ToTensor()
        ])
        self.data = torchvision.datasets.CIFAR10(root='./cifar_data', train=True, 
                                                   download=True, transform=self.transform)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        img, target = self.data[index]
        return img

unet = Unet(
    dim = 128,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = 3,
    layer_attns = (False, True, True, True),
    layer_cross_attns = (False, True, True, True)
).cuda()

imagen = Imagen(
    condition_on_text = False,  # this must be set to False for unconditional Imagen
    unets = unet,
    image_sizes = 32,
    timesteps = 1000
).cuda()

dataset = CifarDataset()

# working training loop
print('starting training loop for 200k iterations')

learning_rate = 1e-4
opt = optim.Adam(unet.parameters(), lr=learning_rate, betas=(0.9, 0.999), eps=1e-08,
                     weight_decay=False)

dataloader = get_images_dataloader(batch_size = 32)
for i in range(200000):
    for image_batch in tqdm(dataloader):
        image_batch = image_batch.cuda()
        loss = imagen(image_batch, unet=unet, unet_number = 1)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)
        opt.step()
        opt.zero_grad()
    print(f'loss: {loss}')
    images = imagen.sample(batch_size = 1, return_pil_images = True) # returns List[Image]
    images[0].save(f'./imagen_samples/sample-{i}.png')