How to use torchvision.models.feature_extraction.create_feature_extractor() with vit_pytorch?

[ArturasDruteika]

I want to create a heatmap for the vit models. To extract necessary layers output I am using torchvision.models.feature_extraction.create_feature_extractor() (detailed explanation https://pytorch.org/blog/FX-feature-extraction-torchvision/).
No matter which model I choose from vit_pytorch library and when using create_feature_extractor(model, return_nodes=return_nodes) I always get the same error: RuntimeError: 'len' is not supported in symbolic tracing by default. If you want this call to be recorded, please call torch.fx.wrap('len') at module scope
I tried adding torch.fx.wrap('len') in various places but I cannot seem to fix this error. How could I solve it?