CDAWG:

        model.subspaces = nn.Parameter(orthogonalization(model.subspaces))


This is a bit dangerous as Parameters need to always be leafs and this might not be the case.
If you just want to update the content of a parameter without the autograd tracking it, you can do:
with torch.no_grad():
    model.subspaces.copy_(orthogonalization(model.subspaces))