[blamestross]

It uses pytorch (and I'm probably going to use it), but doesn't effectively leverage a GPU for training.

[jawarner]

What do you think is the best way to accomplish this?

[blamestross]

You don't. You need a different parallelism model than a GPU provides. It could work well on machines with very high CPU count, but the speedup on GPUs is the main reason bipartite graph algorithms have seen such investment.