# Path: loggers/loggers.py
# class _LogCumulator(BaseLogger):
#     """ cumulate the batch_log and generate an epoch_log
#     Note that this logger is used for generating epoch_log,
#     and thus does not take epoch_log as input"""
# 
#     def __init__(self):
#         pass
# 
#     def on_epoch_begin(self, epoch):
#         self.log_values = defaultdict(list)
#         self.sizes = list()
#         self.epoch_log = None
# 
#     def on_batch_end(self, batch, batch_log):
#         for k, v in batch_log.items():
#             self.log_values[k].append(v)
#         self.sizes.append(batch_log['size'])
# 
#     def get_epoch_log(self):
#         epoch_log = dict()
#         for k in self.log_values:
#             epoch_log[k] = (np.array(self.log_values[k]) * np.array(self.sizes)).sum() / np.array(self.sizes).sum()
#         return epoch_log
# 
# Path: util/util_print.py
# class bcolors:
#     HEADER = '\033[95m'
#     OKBLUE = '\033[94m'
#     OKGREEN = '\033[92m'
#     WARNING = '\033[93m'
#     FAIL = '\033[91m'
#     ENDC = '\033[0m'
#     BOLD = '\033[1m'
#     UNDERLINE = '\033[4m'


import time
import torch.optim as optim
import torch
import torch.distributed as dist
from torch.nn import init
from torch import FloatTensor, tensor
from loggers.loggers import _LogCumulator
from util.util_print import str_warning
        # for v in self._moveable_vars:
        #     if isinstance(v, str):
        #         if '.' in v:
        #             var_type, var_name = v.split('.')
        #             var = getattr(getattr(self, var_type), var_name)
        #             setattr(getattr(self, var_type), var_name, var.to(device))
        #         else:
        #             setattr(self, v, getattr(self, v).to(device))
        #     else:
        #         v.to(device)
        for net in self._nets:
            net.to(device)
        self.device = device

    def save_state_dict(self, filepath, *, save_optimizer=False, additional_values={}):
        state_dicts = dict()
        state_dicts['nets'] = [net.state_dict() for net in self._nets]
        if save_optimizer:
            state_dicts['optimizers'] = [optimizer.state_dict()
                                         for optimizer in self._optimizers]
        for k, v in additional_values.items():
            state_dicts[k] = v
        torch.save(state_dicts, filepath)

    def load_state_dict(self, filepath, *, load_optimizer='auto'):
        state_dicts = torch.load(filepath)

        if load_optimizer == 'auto':
            load_optimizer = ('optimizers' in state_dicts)
            if not load_optimizer:

# Path: loggers/loggers.py
# class _LogCumulator(BaseLogger):
#     """ cumulate the batch_log and generate an epoch_log
#     Note that this logger is used for generating epoch_log,
#     and thus does not take epoch_log as input"""
# 
#     def __init__(self):
#         pass
# 
#     def on_epoch_begin(self, epoch):
#         self.log_values = defaultdict(list)
#         self.sizes = list()
#         self.epoch_log = None
# 
#     def on_batch_end(self, batch, batch_log):
#         for k, v in batch_log.items():
#             self.log_values[k].append(v)
#         self.sizes.append(batch_log['size'])
# 
#     def get_epoch_log(self):
#         epoch_log = dict()
#         for k in self.log_values:
#             epoch_log[k] = (np.array(self.log_values[k]) * np.array(self.sizes)).sum() / np.array(self.sizes).sum()
#         return epoch_log
# 
# Path: util/util_print.py
# class bcolors:
#     HEADER = '\033[95m'
#     OKBLUE = '\033[94m'
#     OKGREEN = '\033[92m'
#     WARNING = '\033[93m'
#     FAIL = '\033[91m'
#     ENDC = '\033[0m'
#     BOLD = '\033[1m'
#     UNDERLINE = '\033[4m'


# Path: models/netinterface.py
import time
import torch.optim as optim
import torch
import torch.distributed as dist
from torch.nn import init
from torch import FloatTensor, tensor
from loggers.loggers import _LogCumulator
from util.util_print import str_warning

        # for v in self._moveable_vars:
        #     if isinstance(v, str):
        #         if '.' in v:
        #             var_type, var_name = v.split('.')
        #             var = getattr(getattr(self, var_type), var_name)
        #             setattr(getattr(self, var_type), var_name, var.to(device))
        #         else:
        #             setattr(self, v, getattr(self, v).to(device))
        #     else:
        #         v.to(device)
        for net in self._nets:
            net.to(device)
        self.device = device

    def save_state_dict(self, filepath, *, save_optimizer=False, additional_values={}):
        state_dicts = dict()
        state_dicts['nets'] = [net.state_dict() for net in self._nets]
        if save_optimizer:
            state_dicts['optimizers'] = [optimizer.state_dict()
                                         for optimizer in self._optimizers]
        for k, v in additional_values.items():
            state_dicts[k] = v
        torch.save(state_dicts, filepath)

    def load_state_dict(self, filepath, *, load_optimizer='auto'):
        state_dicts = torch.load(filepath)

        if load_optimizer == 'auto':
            load_optimizer = ('optimizers' in state_dicts)
            if not load_optimizer:

                print(str_warning, 'Model loaded without optimizer states. ')