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)
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)