@benbenwt please see: UnsupportedOperationException while training and let me know if you have any issues.