import torch
import os
import torch.nn.functional as F

def save_checkpoint(model, filelocation, save_parallel = True):
    if save_parallel:
        torch.save(model.module.state_dict(), filelocation)
    else:
        torch.save(model.state_dict(), filelocation)

def load_Checkpoint(fileLocation,model, load_cpu=False):
    if load_cpu:
        model.load_state_dict(torch.load(fileLocation,map_location=lambda storage, loc: storage))
    else:
        model.load_state_dict(torch.load(fileLocation))
    return model

def writeLog(logList, filename):
    with open(filename, 'w') as outfile:
        outfile.write("\n".join(logList))


def kl_loss(mu, logvar):
    return -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).mean()