示例#1
0
    def save(self, optimizer):
        # set all config to model
        model = self.model
        if self.use_multi_gpu:
            model = self.model.module

        model.train_counter = self.train_counter
        model.metrics = self.metric_logs

        utils.save_checkpoint(self.log_dir, model, optimizer)
示例#2
0
文件: trainer.py 项目: dw-park/claf
    def save(self, optimizer):
        # set all config to model
        model = self.model
        if self.use_multi_gpu:
            model = self.model.module

        model.train_counter = self.train_counter
        model.metrics = self.metric_logs

        if nsml.IS_ON_NSML:
            nsml.save(self.train_counter.get_display())
        else:
            utils.save_checkpoint(self.log_dir, model, optimizer)