learn_sigma must be True?

[MordehayM]

According to this line:

RePaint/guided_diffusion/gaussian_diffusion.py

Line 239 in 8e2a055

assert model_output.shape == (B, C * 2, *x.shape[2:])

I think the lean_sigma argument must be True.
If that is true, please mention this!

Thank

[wdcww]

Hi, I think this may be related to the pre-trained model used. (I used celeba256_250000.pt)
When I change learn_sigma in .yml from True to False, it will prompt:

RuntimeError: Error(s) in loading state_dict for UNetModel:
	size mismatch for out.2.weight: copying a param with shape torch.Size([6, 256, 3, 3]) from checkpoint, the shape in current model is torch.Size([3, 256, 3, 3]).
	size mismatch for out.2.bias: copying a param with shape torch.Size([6]) from checkpoint, the shape in current model is torch.Size([3]).

And I noticed that in the create_model function in ./guided_diffusion/script_util.py, where the value is
returned , It is as follows:

return UNetModel(
        image_size=image_size,
        in_channels=3,
        model_channels=num_channels,
        out_channels=(3 if not learn_sigma else 6),  ## here!
        num_res_blocks=num_res_blocks,
        ... ...