create_feature_extractor problem with VIT

[redeemme]

When I use the create_feature_extractor to get the nodes of VIT,I found errors as follows:

Details

Building Model....Traceback (most recent call last):
File "/home/user5/FGVC-PIM/main_PAC.py", line 350, in
main(args, tlogger)
File "/home/user5/FGVC-PIM/main_PAC.py", line 260, in main
train_loader, val_loader, model, optimizer, schedule, scaler, amp_context, start_epoch = set_environment(args, tlogger)
File "/home/user5/FGVC-PIM/main_PAC.py", line 82, in set_environment
use_combiner = args.use_combiner,
File "/home/user5/FGVC-PIM/models/builder.py", line 222, in build_vit16
comb_proj_size = comb_proj_size)
File "/home/user5/FGVC-PIM/models/pim_module/pim_module_vit.py", line 439, in init
self.backbone = create_feature_extractor(backbone, return_nodes = return_nodes)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torchvision/models/feature_extraction.py", line 441, in create_feature_extractor
graph = tracer.trace(model)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/fx/symbolic_trace.py", line 571, in trace
self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
File "/home/user5/FGVC-PIM/timm/models/vision_transformer.py", line 339, in forward
x = self.forward_features(x)
File "/home/user5/FGVC-PIM/timm/models/vision_transformer.py", line 330, in forward_features
x = self.blocks(x)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/fx/symbolic_trace.py", line 560, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torchvision/models/feature_extraction.py", line 79, in call_module
out = forward(*args, **kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/fx/symbolic_trace.py", line 556, in forward
return _orig_module_call(mod, *args, **kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
return forward_call(*input, **kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/container.py", line 139, in forward
input = module(input)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/fx/symbolic_trace.py", line 560, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torchvision/models/feature_extraction.py", line 79, in call_module
out = forward(*args, **kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/fx/symbolic_trace.py", line 556, in forward
return _orig_module_call(mod, *args, **kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
return forward_call(*input, **kwargs)
File "/home/user5/FGVC-PIM/timm/models/vision_transformer.py", line 206, in forward
x = x + self.drop_path(self.attn(self.norm1(x)))
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/fx/symbolic_trace.py", line 560, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torchvision/models/feature_extraction.py", line 79, in call_module
out = forward(*args, **kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/fx/symbolic_trace.py", line 556, in forward
return _orig_module_call(mod, *args, **kwargs)
File "/home/user5/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
return forward_call(*input, **kwargs)
File "/home/user5/FGVC-PIM/timm/models/vision_transformer.py", line 182, in forward
attn = (q @ k.transpose(-2, -1)) * self.scale
TypeError: unsupported operand type(s) for @: 'Proxy' and 'Proxy'