How can I manually update some of their checkpoints of UNet2/3DConditionModel objects?

^{Originally posted by **justin4ai** August 26, 2024}
Hello, I'm quite new to diffusers package and trying to implement fine-tuning code that uses the saved checkpoints initialized with ```UNet2/3DConditionModel.from_pretrained``` method as shown below:

```python

reference_unet = UNet2DConditionModel.from_pretrained( # ReferenceNet은 2D condition만 받음 (reference image via CLIP)
cfg.base_model_path,
subfolder="unet",
).to(device="cuda")

denoising_unet = UNet3DConditionModel.from_pretrained_2d(
cfg.base_model_path,
"",
subfolder="unet",
unet_additional_kwargs={
"use_motion_module": False,
"unet_use_temporal_attention": False,
},
).to(device="cuda")

prev = denoising_unet.state_dict()

li = torch.load("./pretrained_weights/denoising_unet.pth")

for key in li:
denoising_unet[key] = li[key] # I know this kind of direct assigning to the object doesn't make sense though.
reference_unet.load_state_dict(torch.load("./pretrained_weights/reference_unet.pth"))

```

The checkpoint I try to load is saved from the previous training of ``` UNet2/3DConditionModel objects``` with ```state_dict = model.state_dict()``` and ```torch.save(state_dict, save_path)```. But I have no Idea about how to directly assign certain values to specific layers in those class objects.

If you help me out with this, I will be so much glad! Looking forward to your help. Also please let me know if my description of the situation is not enough for you to help me out.

Cheers,
Junstin