Get prediction working with data models (e.g., LSTMStateTransitionModel) ## DoD - [ ] Implement example - [ ] Make any changes required to get it working - [ ] Add documentation to help users - [ ] Add to release notes