Issues with new LayoutLMv2 
I am working with the newly added LayoutLMv2. Works fine with performing a forward pass, but get a dimensionality error related to the embeddings when I try to use it in another library, namely Captum for explainability. Note that LayoutLM gives no issues in the same context. Also, I realize that this model needs to be finetuned. This is just supposed to be a proof-of-concept usage. 

Here is my code:
```
from PIL import Image, ImageDraw, ImageFont
from transformers import LayoutLMv2FeatureExtractor, LayoutLMv2TokenizerFast, LayoutLMv2Processor, LayoutLMv2ForSequenceClassification
from captum.attr._utils.input_layer_wrapper import ModelInputWrapper
from captum.attr import LayerIntegratedGradients, TokenReferenceBase
import torch
import torchvision
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

image_rgb = Image.open("IMAGE.jpg").convert("RGB")

processor = LayoutLMv2Processor.from_pretrained('microsoft/layoutlmv2-base-uncased')
model = LayoutLMv2ForSequenceClassification.from_pretrained('microsoft/layoutlmv2-base-uncased')
tokenizer = LayoutLMv2TokenizerFast.from_pretrained("microsoft/layoutlmv2-base-uncased")

encoding = processor(image_rgb, return_tensors="pt")

input_ids = encoding['input_ids']
token_type_ids = encoding['token_type_ids']
attention_mask = encoding['attention_mask']
bbox = encoding['bbox']

model_layered = ModelInputWrapper(model)

outputs = model_layered(**encoding)

pred, answer_idx = F.softmax(outputs.logits, dim=1).data.cpu().max(dim=1)

def batch_predict(input_ids, image, bbox, attention_mask, token_type_ids):
    model_layered.eval()
    outputs = model_layered(input_ids=input_ids,
                            image=image,
                            bbox=bbox,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids)
    logits = outputs.logits
    probs = F.softmax(logits, dim=1)
    return probs

attr = LayerIntegratedGradients(batch_predict, 
                                [model_layered.module.layoutlmv2.embeddings.word_embeddings, 
                                 model_layered.module.layoutlmv2.embeddings.position_embeddings, 
                                 model_layered.module.layoutlmv2.embeddings.x_position_embeddings,
                                 model_layered.module.layoutlmv2.embeddings.y_position_embeddings,
                                 model_layered.module.layoutlmv2.embeddings.h_position_embeddings, 
                                 model_layered.module.layoutlmv2.embeddings.w_position_embeddings,
                                 model_layered.module.layoutlmv2.embeddings.token_type_embeddings,])

# Generate reference for tokens
token_reference = TokenReferenceBase(reference_token_idx=tokenizer.pad_token_id)
text_reference_indices = token_reference.generate_reference(len(encoding['input_ids'][0]), device=device).unsqueeze(0)
baselines = text_reference_indices

attributions = attr.attribute(inputs=encoding['input_ids'],
                              additional_forward_args=(encoding['image'], 
                                                       encoding['bbox'], 
                                                       encoding['attention_mask'],
                                                       encoding['token_type_ids']),
                            baselines=baselines,
                            target=answer_idx,
                            n_steps=5)
```

And the error: 
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-1-00bdf4f97de3> in <module>()
     60                             baselines=baselines,
     61                             target=answer_idx,
---> 62                             n_steps=5)

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/captum/log/__init__.py in wrapper(*args, **kwargs)
     33             @wraps(func)
     34             def wrapper(*args, **kwargs):
---> 35                 return func(*args, **kwargs)
     36 
     37             return wrapper

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/captum/attr/_core/layer/layer_integrated_gradients.py in attribute(self, inputs, baselines, target, additional_forward_args, n_steps, method, internal_batch_size, return_convergence_delta, attribute_to_layer_input)
    496             method=method,
    497             internal_batch_size=internal_batch_size,
--> 498             return_convergence_delta=False,
    499         )
    500 

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/captum/attr/_core/integrated_gradients.py in attribute(self, inputs, baselines, target, additional_forward_args, n_steps, method, internal_batch_size, return_convergence_delta)
    290                 additional_forward_args=additional_forward_args,
    291                 n_steps=n_steps,
--> 292                 method=method,
    293             )
    294 

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/captum/attr/_core/integrated_gradients.py in _attribute(self, inputs, baselines, target, additional_forward_args, n_steps, method, step_sizes_and_alphas)
    353             inputs=scaled_features_tpl,
    354             target_ind=expanded_target,
--> 355             additional_forward_args=input_additional_args,
    356         )
    357 

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/captum/attr/_core/layer/layer_integrated_gradients.py in gradient_func(forward_fn, inputs, target_ind, additional_forward_args)
    464 
    465                     output = _run_forward(
--> 466                         self.forward_func, tuple(), target_ind, additional_forward_args
    467                     )
    468                 finally:

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/captum/_utils/common.py in _run_forward(forward_func, inputs, target, additional_forward_args)
    451         *(*inputs, *additional_forward_args)
    452         if additional_forward_args is not None
--> 453         else inputs
    454     )
    455     return _select_targets(output, target)

<ipython-input-1-00bdf4f97de3> in batch_predict(input_ids, image, bbox, attention_mask, token_type_ids)
     34                             bbox=bbox,
     35                             attention_mask=attention_mask,
---> 36                             token_type_ids=token_type_ids)
     37     logits = outputs.logits
     38     probs = F.softmax(logits, dim=1)

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/captum/attr/_utils/input_layer_wrapper.py in forward(self, *args, **kwargs)
     74             kwargs[arg_name] = self.input_maps[arg_name](kwargs[arg_name])
     75 
---> 76         return self.module(*tuple(args), **kwargs)

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/transformers/models/layoutlmv2/modeling_layoutlmv2.py in forward(self, input_ids, bbox, image, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, labels, output_attentions, output_hidden_states, return_dict)
   1053             output_attentions=output_attentions,
   1054             output_hidden_states=output_hidden_states,
-> 1055             return_dict=return_dict,
   1056         )
   1057         if input_ids is not None:

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/transformers/models/layoutlmv2/modeling_layoutlmv2.py in forward(self, input_ids, bbox, image, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict)
    893             token_type_ids=token_type_ids,
    894             position_ids=position_ids,
--> 895             inputs_embeds=inputs_embeds,
    896         )
    897 

/home/natbarkas/anaconda3/envs/explainability-v2/lib/python3.6/site-packages/transformers/models/layoutlmv2/modeling_layoutlmv2.py in _calc_text_embeddings(self, input_ids, bbox, position_ids, token_type_ids, inputs_embeds)
    754         token_type_embeddings = self.embeddings.token_type_embeddings(token_type_ids)
    755 
--> 756         embeddings = inputs_embeds + position_embeddings + spatial_position_embeddings + token_type_embeddings
    757         embeddings = self.embeddings.LayerNorm(embeddings)
    758         embeddings = self.embeddings.dropout(embeddings)

RuntimeError: The size of tensor a (44) must match the size of tensor b (49) at non-singleton dimension 1
```
Is it possible to make a Colab notebook that can reproduce the issue?<|||||>Yes, I sent you a link to the Colab notebook. 