"unexpected keyword argument 'head_fusion'" and "required positional argument: 'category_index'"

[expectopatronum]

Hi,
I tried to run the example from the README.md and I run into the following issues:

• grad_rollout = VITAttentionGradRollout(model, discard_ratio=0.9, head_fusion='max') causes TypeError: __init__() got an unexpected keyword argument 'head_fusion'.
• mask = grad_rollout(input_tensor) causes TypeError: __call__() missing 1 required positional argument: 'category_index' although the README suggests it is optional (If category_index isn't specified, Attention Rollout will be used, otherwise Gradient Attention Rollout will be used.). Is this a mistake or should I pass None for example?

Thanks a lot and best regards
Verena