Thanks for the great library.
There seems to be an incompatibility issue with the latest torch.func and transformers as the continual learning example no longer works. The error I got was RuntimeError: vmap: inplace arithmetic(self, *extra_args) is not possible because there exists a Tensor other in extra_args that has more elements than self.
Have you seen this error and if yes, do you have a workaround, and if no, could you share the versions of pytorch and transformers used for the public example?