In the proposed solution on GitHub, the first solution gives you a different dataset in each epoch. The second solution will give you a batch of both the datasets in each step for all epochs.
With the second solution, you can simply calculate the loss by computing it separately with each dataset. With the first solution I don’t see how that’s possible. Maybe if you can share the exact workflow of your use-case, I can suggest a solution then.