I suggest a simple MSE or BCE loss will do.

I tried, but it keeps giving me a huge negative number like below.
Do you have any idea what is actually going on?
Input to VAE is a range of[0 - 1] and output has the same range of number too.
And also, MSE didn’t work.
Reconstructed data is not reproducing input data at all.
With changing the param for KLD, I could get better results.
But I don’t want to take this approach since it won’t be a general solution anymore if I changed the input data.
 41%|████      | 13/32 [01:04<01:27,  4.61s/it]
kld :  53.88439178466797 recon :  -53301.86328125
 44%|████▍     | 14/32 [01:08<01:22,  4.58s/it]
kld :  52.833587646484375 recon :  -56408.796875
 47%|████▋     | 15/32 [01:13<01:18,  4.59s/it]
kld :  51.79539489746094 recon :  -59532.7265625
 50%|█████     | 16/32 [01:17<01:12,  4.56s/it]
kld :  50.653541564941406 recon :  -62645.80859375
 53%|█████▎    | 17/32 [01:22<01:08,  4.54s/it]
kld :  49.04255676269531 recon :  -65793.15625
 56%|█████▋    | 18/32 [01:26<01:03,  4.52s/it]
kld :  47.729461669921875 recon :  -68937.03125
 59%|█████▉    | 19/32 [01:31<00:58,  4.51s/it]
kld :  47.20650100708008 recon :  -72088.6171875
 62%|██████▎   | 20/32 [01:35<00:54,  4.51s/it]
kld :  46.273780822753906 recon :  -75279.53125
 66%|██████▌   | 21/32 [01:40<00:49,  4.51s/it]
kld :  44.844871520996094 recon :  -78400.25
 69%|██████▉   | 22/32 [01:44<00:44,  4.49s/it]
kld :  42.69890213012695 recon :  -81610.5
・
・
・

What do you think about the function below?
Do you have any idea?

inputs.shape = (512, 19, 64, 64)
outputs.shape = (512, 19, 64, 64)

def entropy(inputs, outputs, weight,batch=batch_size, h=64, w=64, class_num=19, ignore_id=19):
    
    softmax = nn.Softmax(dim=0)
    losses = torch.empty(h, w).to(device)
    
    for b in range(batch):
        
        loss = - softmax(inputs[b]) * torch.log(softmax(outputs[b]))
        
        for c in range(class_num):
            if c == ignore_id:
                continue
                
            weighted_loss = ((weight[c]) * loss[c]) / batch_size
            torch.cat((losses, weighted_loss))
            
    return torch.sum(losses)