.predict_embedding() for DMPNNModel (TorchModel)

Hi, I am new to Deepchem and I am experimenting with the DMPNNModel from PyTorch models. I would like to use .predict_embedding() to analyze the embedding space, but it does seem to work for me (not yet implemented for DMPNNModel class?). I would deeply appreciate any advise, insights, ideas and how this can be accomplished, or perhaps suggestions for alternative approaches?

Many thanks in advance,
Dmitriy