We extend the DeePMD neural network architecture to predict electronic structure properties necessary to perform non-adiabatic dynamics simulations. While learning the excited state energies and forces follows a straightforward extension of the DeePMD approach for ground-state energies and forces, how to learn the map between the non-adiabatic coupling vectors (NACV) and the local chemical environment descriptors of DeePMD is less trivial. Most implementations of machine-learning-based non-adiabatic dynamics inherently approximate the NACVs, with an underlying assumption that the energy-difference-scaled NACVs are conservative fields. We overcome this approximation, implementing the method recently introduced by Richardson [J. Chem. Phys. 158 011102 (2023)], which learns the symmetric dyad of the energy-difference-scaled NACV. The efficiency and accuracy of our neural network architecture is demonstrated through the example of the methaniminium cation CH$_2$NH$_2^+$.
