fix prior saving / add load method about dalle2-pytorch HOT 3 CLOSED

from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork
import torch

device = "cuda:0"
state_dict = torch.load("model.pth")

prior_network = DiffusionPriorNetwork( 
    dim=768,
    depth=6, 
    dim_head=64, 
    heads=8
).to(device)

diffusion_prior = DiffusionPrior( 
    net=prior_network,
    clip=None, 
    image_embed_dim=768, 
    timesteps=100,
    cond_drop_prob=0.2, 
    loss_type="l2", 
    condition_on_text_encodings=False
).to(device)

diffusion_prior.load_state_dict(state_dict)