def log_cm(self, tag, cm, labels=None, step=None):
    step = (step or self.global_step)
    fig = utils.draw_confusion_matrix(cm.cpu())
    image_tensor = utils.fig2rgb(fig)
    self.logger.experiment.add_image(tag, image_tensor, step)
    mlflow.log_figure(fig, (tag + f'/{step}.png'))
def log_cm(self, tag, cm, labels=None, step=None):
    step = (step or self.global_step)
    fig = utils.draw_confusion_matrix(cm.cpu())
    image_tensor = utils.fig2rgb(fig)
    self.logger.experiment.add_image(tag, image_tensor, step)
    mlflow.log_figure(fig, (tag + f'/{step}.png'))