Error loading weights for maze_2d

[Abc11c]

Hi,

I was trying to use the weights from #13 but I'm getting the following error when I run python scripts/plan_maze_2d.py

Traceback (most recent call last):
  File "scripts/plan_maze2d.py", line 25, in <module>
    diffusion_experiment = utils.load_diffusion(args.logbase, args.dataset, args.diffusion_loadpath, epoch=args.diffusion_epoch)
  File "/home/karthi43/working/diffuser/diffuser/utils/serialization.py", line 58, in load_diffusion
    trainer.load(epoch)
  File "/home/karthi43/working/diffuser/diffuser/utils/training.py", line 162, in load
    self.model.load_state_dict(data['model'])
  File "/home/karthi43/miniconda3/envs/diffuser/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1406, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for GaussianDiffusion:
        Missing key(s) in state_dict: "model.downs.0.3.conv.weight", "model.downs.0.3.conv.bias", "model.downs.1.3.conv.weight", "model.downs.1.3.conv.bias", "model.ups.0.3.conv.weight", "model.ups.0.3.conv.bias", "model.ups.1.3.conv.weight", "model.ups.1.3.conv.bias". 
        Unexpected key(s) in state_dict: "model.downs.0.2.conv.weight", "model.downs.0.2.conv.bias", "model.downs.1.2.conv.weight", "model.downs.1.2.conv.bias", "model.ups.0.2.conv.weight", "model.ups.0.2.conv.bias", "model.ups.1.2.conv.weight", "model.ups.1.2.conv.bias". 

I'm on maze_2d branch, am I missing some configs or something

Thanks you!

[Abc11c]

Works when I re-train stuff on my end and then perform inference