Thank you for the feedback. I end up using something like:
    def flatten_params(self):
        l = [torch.flatten(p) for p in self.parameters()]
        indices = []
        s = 0
        for p in l:
            size = p.shape[0]
            indices.append((s, s+size))
            s += size
        flat = torch.cat(l).view(-1, 1)
        return {"params": flat, "indices": indices}

    def recover_flattened(self, flat_params, indices):
        l = [flat_params[s:e] for (s, e) in indices]
        for i, p in enumerate(self.parameters()):
            l[i] = l[i].view(*p.shape)
            p.data = l[i]