Make the models work on multiple GPUs

[Clifford-appliedAI-GmbH]

Last time I checked we were getting some data related errors. It's possible that the latest Lightning related changes in #33 have helped to alleviate this.