This won’t be possible since the modelX are already detached by the threshold operation as well as the argmin.
You would thus have to use a differentiable / “soft” threshold operation as described e.g. here by @KFrank.
1 Like