perf(aggregation): Prevent cuda sync in normalize

[ValerianRey]

With the previous implementation, normalize causes a cuda synchronization. It's not a big issue at all, because we need a cuda synchronization right after, during the call to project_weights, because it's always done on CPU. But I think it's good for three reasons:

• Keep the number of cuda synchronizations minimal for performance reasons.
• Avoiding synchronization in unexpected places makes it easier to analyze traces.
• If we ever do UPGrad on cuda, it could make a big difference to avoid this cuda sync, because we may be down to 0 cuda sync for aggregation.

There might be a slight performance drop due to using torch.where which is element-wise with a condition that is scalar (and thus broadcasted). But since the gramian is never huge (especially if using UPGrad), this is really fine IMO. In my profiling, this torch.where takes 0.028 ms with batch size of 64.

So this is extremely minor but positive IMO.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/torchjd/_linalg/_gramian.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[PierreQuinton]

Very nice, we should also keep that in mind because I doubt this is the only place we do something like that (the aggregators maybe).