@tom
Thanks. It works now
I have one more question to ask.
I want to use this loss combined with MSE error loss to optimize my network
            loss = alpha * loss_mse + (1 - alpha) * loss_kl_div
            error += loss.item()
            loss.backward()
            optimizer.step()

Will this work for what I wish to achieve?