Function load_state_dict

diffportrait360_release/code/train.py:32–37

(model, ckpt_path, strict=True, map_location="cpu")

print(f"TORCH_VERSION={TORCH_VERSION} FP16_DTYPE={FP16_DTYPE}")
torch.set_default_tensor_type(torch.FloatTensor)
def load_state_dict(model, ckpt_path, strict=True, map_location="cpu"):
    print(f"Loading model state dict from {ckpt_path} ...")
    state_dict = load_from_pretrain(ckpt_path, map_location=map_location)
    state_dict = state_dict.get('state_dict', state_dict)
    model.load_state_dict(state_dict, strict=strict)
    del state_dict  

def get_cond_control(args, batch_data, control_type, device, model=None,batch_size=None, train=True, seg_model=None):
    # Single-control