'NoneType' object has no attribute 'opt' Error, for testing

[kitrak-rev]

With the code reference from Link
built the following code in jupyter-notebook

1. Cell-1

!pip install dalle-pytorch==0.14.3
!pip install deepspeed==0.4.2
!pip install triton==0.4.2
!pip3 install Pillow==8.2.0
!pip install triton

2. Cell-2

%matplotlib inline
import base64
from io import BytesIO

import torch
import numpy as np
import matplotlib.pyplot as plt

# vision imports
from PIL import Image
# dalle related classes and utils
from dalle_pytorch import VQGanVAE, DALLE,OpenAIDiscreteVAE
from dalle_pytorch.tokenizer import tokenizer
from einops import repeat 

3. Cell-3

from dalle_pytorch import VQGanVAE
vae = VQGanVAE()
# gan2 = GAN(params)

4. Cell-4

load_obj = torch.load("./16L_64HD_8H_512I_128T_cc12m_cc3m_3E.pt")  # model checkpoint : https://github.com/robvanvolt/DALLE-models/tree/main/models/taming_transformer
dalle_params, _, weights = load_obj.pop('hparams'), load_obj.pop('vae_params'), load_obj.pop('weights')
dalle_params.pop('vae', None)  # cleanup later

dalle = DALLE(vae=vae, **dalle_params).cuda()

dalle.load_state_dict(weights)

batch_size = 4

top_k = 0.9

# generate images

image_size = vae.image_size

5. Cell-5

text = torch.randint(0, 10000, (4, 256)).cuda()
images = dalle.generate_images(text)

When running Cell-5 I get the following error

AttributeError: 'NoneType' object has no attribute 'opt'

Notebook can be replicated at colab link