torch_cg
torch_cg copied to clipboard
Only call A_bmm once per iter, rather than three times
In general, the most expensive part of any CG solve is computing A_bmm
. This version will do it once per iteration rather than three times.