def save(self, path):
    rank = os.environ.get('LOCAL_RANK')
    if ((rank is not None) and (rank != '0')):
        return
    path = utils.makedir(path)
    if self.save_as_state_dict:
        torch.save(self.model.state_dict(), os.path.join(path, 'model_state.pt'))
    else:
        torch.save(self.model, os.path.join(path, 'model.pt'))
    torch.save({'epoch': self.epoch, 'optimizer': self.optimizer.state_dict()}, os.path.join(path, 'optimizer.pt'))
    if (self.scheduler is not None):
        torch.save(self.scheduler.state_dict(), os.path.join(path, 'scheduler.pt'))
    self.logger.info('Saved checkpoint to %s', path)
def save(self, path):
    rank = os.environ.get('LOCAL_RANK')
    if ((rank is not None) and (rank != '0')):
        return
    path = utils.makedir(path)
    if self.save_as_state_dict:
        torch.save(self.model.state_dict(), os.path.join(path, 'model_state.pt'))
    else:
        torch.save(self.model, os.path.join(path, 'model.pt'))
    torch.save({'epoch': self.epoch, 'optimizer': self.optimizer.state_dict()}, os.path.join(path, 'optimizer.pt'))
    if (self.scheduler is not None):
        torch.save(self.scheduler.state_dict(), os.path.join(path, 'scheduler.pt'))
    self.logger.info('Saved checkpoint to %s', path)