Training is fine, but i can not load the trained model

[zhrli]

Describe the bug
As I run this code, errors reported:
import torch
from dgmr import DGMR, Sampler, Generator, Discriminator, LatentConditioningStack, ContextConditioningStack

model = DGMR(
forecast_steps=20, #20
input_channels=1,
output_shape=256,
latent_channels=384,
context_channels=192,
num_samples=3,
)

model.sampler = Sampler.from_pretrained("openclimatefix/dgmr-sampler")
model.sampler.forecast_steps = 20
model.sampler.output_shape = 256
model.discriminator = Discriminator.from_pretrained("openclimatefix/dgmr-discriminator")
model.latent_stack = LatentConditioningStack.from_pretrained("openclimatefix/dgmr-latent-conditioning-stack")
model.context_stack = ContextConditioningStack.from_pretrained("openclimatefix/dgmr-context-conditioning-stack")
model.generator = Generator(conditioning_stack=model.context_stack, latent_stack=model.latent_stack, sampler=model.sampler)
model=DGMR.load_from_checkpoint(checkpoint_path='best-v3.ckpt', strict=False)
print(model.config)

To Reproduce
Steps to reproduce the behavior:
RuntimeError: Error(s) in loading state_dict for DGMR:
size mismatch for latent_stack.l_block1.conv_1x1.weight: copying a param with shape torch.Size([16, 8, 1, 1]) from checkpoint, the shape in current model is torch.Size([4, 8, 1, 1]).
...
...
...
Process finished with exit code 1

Expected behavior
Training is fine . but I don't know why i can not load the trained model.

Additional context
Add any other context about the problem here.