def on_test_end(self):
    tmp_path = 'outputs/test_predictions.csv'
    self.trainer.datamodule.val_history['test'].to_csv(tmp_path)
    mlflow.log_artifact(tmp_path, 'test')
def on_test_end(self):
    tmp_path = 'outputs/test_predictions.csv'
    self.trainer.datamodule.val_history['test'].to_csv(tmp_path)
    mlflow.log_artifact(tmp_path, 'test')