Question about dynamic merging

[mynameisissac]

Thanks for the interesting work.

Regarding the dynamic merging, why should average the predicted coeffs over the batch dimension? Although this gives higher efficiency, shouldn't the dynamic merging be 'sample-wise'?
In such case, I think the sample order will affect the prediction? or also other factors involving the batch dimension will have effect.

mass/src/mass/pl_module/mass.py

Lines 143 to 147 in 583b35e

	merged_model = self._apply_tv(
	list(dataset_group),
	coefficients=dataset_coeffs[
	assigned_sample_idxs[:, None], dataset_group_idxs
	].mean(dim=0),

[crisostomi]

Hello and thank you for the issue!

You are right that they should be kept sample-wise. Indeed, we don't actually use those coefficients in the end as you can see here

So the error doesn't influence the results in the paper, but it's still something to fix!
we are currently refactoring the codebase so we'll take care of it ASAP