Weights mismatch error. Is my Model Depth wrong?

[TomorrowToday]

I am having some issues generating from my checkpoints. I suspect my Model Depth may need adjustment, but am not sure enough to do it and waste a lot of training time. I have trained a DALL-E model with the included pre-trained VQGanVAE as the vae. My settings for DALL-E Training and Generation are as below:

#--------------------
#Training Settings
#--------------------
EPOCHS = 10
BATCH_SIZE = 4
LEARNING_RATE = 3e-4
LR_DECAY = True
GRAD_CLIP_NORM = 0.5

#--------------------
# Model Settings
#--------------------
MODEL_DIM = 512
TEXT_SEQ_LEN = 512
DEPTH = 2
HEADS = 8
DIM_HEAD = 64
REVERSIBLE = False
LOSS_IMG_WEIGHT = 7
FF_DROPOUT = 0.0
ATTN_DROPOUT = 0.0
STABLE = False
SHIFT_TOKENS = False
ROTARY_EMB = False
ATTN_TYPES = ("full")

#--------------------
# Generate Settings
#--------------------
NUM_IMAGES = 128
BATCH_SIZE = 4
TOP_K = 0.9
OUTPUTS_DIR = './outputs'
BPE_PATH = None
HUG = False
CHINESE = False
TAMING = True
GENTEXT = False

Attempts to generate result in an Error when weights are loaded by dalle:

RuntimeError: Error(s) in loading state_dict for DALLE:
	Missing key(s) in state_dict: "vae.codebook.weight", "vae.encoder.0.0.weight", "vae.encoder.0.0.bias", "vae.encoder.1.0.weight", "vae.encoder.1.0.bias", "vae.encoder.2.0.weight", "vae.encoder.2.0.bias", "vae.encoder.3.net.0.weight", "vae.encoder.3.net.0.bias", "vae.encoder.3.net.2.weight", "vae.encoder.3.net.2.bias", "vae.encoder.3.net.4.weight", "vae.encoder.3.net.4.bias", "vae.encoder.4.net.0.weight", "vae.encoder.4.net.0.bias", "vae.encoder.4.net.2.weight", "vae.encoder.4.net.2.bias", "vae.encoder.4.net.4.weight", "vae.encoder.4.net.4.bias", "vae.encoder.5.weight", "vae.encoder.5.bias", "vae.decoder.0.weight", "vae.decoder.0.bias", "vae.decoder.1.net.0.weight", "vae.decoder.1.net.0.bias", "vae.decoder.1.net.2.weight", "vae.decoder.1.net.2.bias", "vae.decoder.1.net.4.weight", "vae.decoder.1.net.4.bias", "vae.decoder.2.net.0.weight", "vae.decoder.2.net.0.bias", "vae.decoder.2.net.2.weight", "vae.decoder.2.net.2.bias", "vae.decoder.2.net.4.weight", "vae.decoder.2.net.4.bias", "vae.decoder.3.0.weight", "vae.decoder.3.0.bias", "vae.decoder.4.0.weight", "vae.decoder.4.0.bias", "vae.decoder.5.0.weight", "vae.decoder.5.0.bias", "vae.decoder.6.weight", "vae.decoder.6.bias". 
	Unexpected key(s) in state_dict: "vae.model.encoder.conv_in.weight", "vae.model.encoder.conv_in.bias", "vae.model.encoder.down.0.block.0.norm1.weight", "vae.model.encoder.down.0.block.0.norm1.bias", "vae.model.encoder.down.0.block.0.conv1.weight", "vae.model.encoder.down.0.block.0.conv1.bias", "vae.model.encoder.down.0.block.0.norm2.weight", "vae.model.encoder.down.0.block.0.norm2.bias", "vae.model.encoder.down.0.block.0.conv2.weight", "vae.model.encoder.down.0.block.0.conv2.bias", "vae.model.encoder.down.0.block.1.norm1.weight", "vae.model.encoder.down.0.block.1.norm1.bias", "vae.model.encoder.down.0.block.1.conv1.weight", "vae.model.encoder.down.0.block.1.conv1.bias", "vae.model.encoder.down.0.block.1.norm2.weight", "vae.model.encoder.down.0.block.1.norm2.bias", "vae.model.encoder.down.0.block.1.conv2.weight", "vae.model.encoder.down.0.block.1.conv2.bias", "vae.model.encoder.down.0.downsample.conv.weight", "vae.model.encoder.down.0.downsample.conv.bias", "vae.model.encoder.down.1.block.0.norm1.weight", "vae.model.encoder.down.1.block.0.norm1.bias", "vae.model.encoder.down.1.block.0.conv1.weight", "vae.model.encoder.down.1.block.0.conv1.bias", "vae.model.encoder.down.1.block.0.norm2.weight", "vae.model.encoder.down.1.block.0.norm2.bias", "vae.model.encoder.down.1.block.0.conv2.weight", "vae.model.encoder.down.1.block.0.conv2.bias", "vae.model.encoder.down.1.block.1.norm1.weight", "vae.model.encoder.down.1.block.1.norm1.bias", "vae.model.encoder.down.1.block.1.conv1.weigh...
	size mismatch for image_emb.weight: copying a param with shape torch.Size([1024, 512]) from checkpoint, the shape in current model is torch.Size([8192, 512]).
	size mismatch for to_logits.1.weight: copying a param with shape torch.Size([50944, 512]) from checkpoint, the shape in current model is torch.Size([58112, 512]).
	size mismatch for to_logits.1.bias: copying a param with shape torch.Size([50944]) from checkpoint, the shape in current model is torch.Size([58112]).

I assume I have misconfigured something, but do not understand enough to see what I did wrong.

[rom1504]

You need to use the --taming argument when running generate.py
The error is saying you're trying to use a checkpoint trained with vqgan using VQ-VAE settings

[TomorrowToday]

Taming is set to true, so that is not entirely the problem. It seems I was using some old code to train my dalle model. That code saved vae_params into the dalle checkpoints and now the generate.py relies on vae_params being null in addition to taming being set. The dalle training code, I borrowed from is still being linked to from the readme so it might be worth updating that notebook or noting that it is deprecated.