why does target and model_output have different num of channel?

[ssxfc]

in my settings, target's channel is 1. model_output's channel is 2
and target represents noise
please tell me why have this : terms["mse"] = mean_flat((target - model_output) ** 2)

[JuliaWolleb]

In this setting, you have one output channel for the mean and one for the variance. The line you posted is just the MSE loss used for training of the U-Net, which predicts the noise schedule epsilon.

[ssxfc]

By the way, in my settings, learn_stigma is set to False. May I ask if I should change the dual channel output of the model to a single channel. Or can I simply modify "terms["mse"] = mean_flat((target - model_output) ** 2)" to "terms["mse"] = mean_flat((target - model_output[:, 0, ...]) ** 2)" ?