Thank you for the reply!
I’ve implemented the suggestion you made and it all appears to be working. With your comment and cross-referencing with the docs, it makes much more sense on what’s happening now. Thank you also for the additional comments, I didn’t realize the to(...) had that problem and I made that change.
I do agree that if the docs listed it as the output tensors that this would make more sense for newcomers in the future.