Add step info dictionary to MLflowLoggerCallback with Tune

RLlib
4

@fedetask,

I do not think that is quite how it should work, but I could be wrong. In my understanding, the MLflowLoggerCallback is a tune callback and the CustomCallback is an rllib callback. I think it should look something like this.

tune.run(config={...,  "callbacks": CustomCallback,},
         callbacks=[MLflowLoggerCallback()])
2 Likes