Is there any way that I can use the HuggingFace Transformers as Pyro models?
Author: h56cho
Pull request: no

Hello,

`Pyro` is a Python module that allows users to convert a given (frequentist) neural network into a Bayesian neural network.
I can convert a HuggingFace Transformer into a Pyro model like below:
```python

import torch
from torch import distributions
from transformers import RobertaTokenizer, RobertaForMultipleChoice
import pyro
import pyro.infer
import pyro.optim
import pyro.distributions as dist
import pyro.nn.module as module
from torch import nn
from pyro.infer import SVI

# get the pre-trained HuggingFace RobertaForMultipleChoice and resize the token embeddings after adding the special token
model_RobertaForMultipleChoice = RobertaForMultipleChoice.from_pretrained('roberta-large', output_hidden_states = True)
        
module.to_pyro_module_(model_RobertaForMultipleChoice)
        
# Now we can attempt to be fully Bayesian:
for m in model_RobertaForMultipleChoice.modules():
      for name, value in list(m.named_parameters(recurse=False)):
                setattr(m, name, module.PyroSample(prior=dist.Normal(0, 1)
                                              .expand(value.shape)
                                              .to_event(value.dim())))

# define parameters for training      
guide_delta = guides.AutoDelta(model_RobertaForMultipleChoice)
```

But when I try to compute the mc_loss from this Bayesian Transformer, Python generates an error:

```python
mc_loss = model(input_ids = input_ids, attention_mask = attention_mask, labels = mc_labels)[0]

Traceback (most recent call last):
  File "STAT946_final_project_code_v4.py", line 625, in <module>
    success_rate_list_diag_normal = main_function_diag_normal('/home/ec2-user/test.txt', 'test_ans_num.txt', num_iter, log_interval)
  File "STAT946_final_project_code_v4.py", line 415, in main_function_diag_normal
    best_model_RobertaForMultipleChoice_diag_normal = train_loop(model_RobertaForMultipleChoice, tokenizer, optimizer_1, scheduler_1, log_interval, svi_diag_normal, guide_diag_normal, best_model_RobertaForMultipleChoice_diag_normal)
  File "STAT946_final_project_code_v4.py", line 342, in train_loop
    optimizer, scheduler, log_interval, svi, guide, epoch)
  File "STAT946_final_project_code_v4.py", line 237, in train_mc_head
    mc_loss = model(input_ids = input_ids, attention_mask = attention_mask, labels = mc_labels)[0]
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/pyro/nn/module.py", line 413, in __call__
    return super().__call__(*args, **kwargs)
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/transformers/modeling_roberta.py", line 441, in forward
    output_hidden_states=output_hidden_states,
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/pyro/nn/module.py", line 413, in __call__
    return super().__call__(*args, **kwargs)
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/transformers/modeling_bert.py", line 732, in forward
    extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/transformers/modeling_utils.py", line 228, in get_extended_attention_mask
    extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
  File "/home/ec2-user/anaconda3/lib/python3.7/site-packages/transformers/modeling_utils.py", line 159, in dtype
    first_tuple = next(gen)
StopIteration
```

Is there any way that I can compute the mc_loss in the regular way after converting HuggingFace Transformer into a Bayesian Transformer?

Thank you.,

Comments: