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)