AttributeError: 'MlflowClient' object has no attribute 'add_figure'

[polkadot21]

• PyTorch-Forecasting version: 0.10.1
• PyTorch version: 1.11.0+cu113
• Python version: 3.7.13
• Operating System: VERSION="18.04.5 LTS (Bionic Beaver)" VERSION_ID="18.04" VERSION_CODENAME=bionic

Expected behavior

I executed code to fit my TFT model with the MLFlow logger in order to create an MLFlow run and expected to get a new folder created with hyperparameters and metrics and etc.

Actual behavior

However, result was AttributeError: 'MlflowClient' object has no attribute 'add_figure' I think it has to do with the module base_model.py because of self.logger.experiment.add_figure(...)

Code to reproduce the problem

# configure network and trainer
early_stop_callback = EarlyStopping(monitor="val_loss",
                                    min_delta=1e-4,
                                    patience=5,
                                    verbose=True,
                                    mode="min")

lr_logger = LearningRateMonitor()  # log the learning rate
logger = MLFlowLogger("lightning_logs")  # logging results to a tensorboard


SWA = StochasticWeightAveraging()


trainer = pl.Trainer(
    max_epochs=100,
    gpus=1,
    accelerator="gpu",
    weights_summary="top",
    gradient_clip_val=0.5,
    callbacks=[lr_logger, early_stop_callback, SWA],
    logger=logger,
)


tft = TemporalFusionTransformer.from_dataset(
    training,
    learning_rate=0.008511380382023764,
    hidden_size=56,
    attention_head_size=2,
    dropout=0.217,
    hidden_continuous_size=27,
    output_size=7,  # 7 quantiles by default
    loss=QuantileLoss(),
    log_interval=10,  # uncomment for learning rate finder and otherwise, e.g. to 10 for logging every 10 batches
    reduce_on_plateau_patience=4,
)
print(f"Number of parameters in network: {tft.size()/1e3:.1f}k")

trainer.fit(
    tft,
    train_dataloaders=train_dataloader,
    val_dataloaders=val_dataloader,
)

Paste the command(s) you ran and the output. Including a link to a colab notebook will speed up issue resolution. If there was a crash, please include the traceback here. The code used to initialize the TimeSeriesDataSet and model should be also included.

https://colab.research.google.com/drive/1_0kDkq11cAqa_zlVcncZ3s9sn-OkvjrM?authuser=1#scrollTo=rxRGQ6EWGLCb