I have trained using this as well but not in Databricks (it's what I'm trying to get working). For some reason, the paths that I'm specifying the model cannot find. What did your paths look like to your datasets?