import numpy as np
import torch
from logger import Logger
LOGGER = Logger().logger()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


#Author: @meetdoshi
def initialize_weights(model):
    '''
    @params
    @return
    '''
    for m in model.modules():
        if isinstance(m, torch.nn.Conv2d):
            torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
        elif isinstance(m, torch.nn.BatchNorm2d):
            torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
            torch.nn.init.constant_(m.bias.data, 0)
        elif isinstance(m, torch.nn.Linear):
            torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
            torch.nn.init.constant_(m.bias.data, 0)
