What kind of error you usually see? Sometimes it is memory problem, is it work with a small toy model?