fix cross entropy bug
@@ -207,8 +207,8 @@ class ImageColorizeModule(RunModule, ImageServing):
out_class, out_reg = self(img['A'], img['hint_B'], img['mask_B'])
# loss
- criterionCE = nn.loss.CrossEntropyLoss()
- loss_ce = criterionCE(out_class, img['real_B_enc'][:, 0, :, :])
+ loss_ce = F.cross_entropy(out_class, img['real_B_enc'][:, :1, :, :], axis=1)
+ loss_ce = paddle.mean(loss_ce)
loss_G_L1_reg = paddle.sum(paddle.abs(img['B'] - out_reg), axis=1, keepdim=True)
loss_G_L1_reg = paddle.mean(loss_G_L1_reg)
loss = loss_ce + loss_G_L1_reg
