def __init__(self, kl_tolerance):
    super(KLLoss, self).__init__()
    self.kl_tolerance = torch.tensor(kl_tolerance)
def __init__(self, kl_tolerance):
    super(KLLoss, self).__init__()
    self.kl_tolerance = torch.tensor(kl_tolerance)