Change model from bert into electra - FloatTensor changes to tuple?

Hello,

I would like to add electra to the wl-coref model.
In doing so, I received a type change in the variable and thus the whole thing does not work.

If the model is bert, then the variable “out” has the type torch.FloatTensor and the dtype torch.float32.
With electra “out” has the type tuple.
The error message that appears: “TypeError: only integer tensors of a single element can be converted to an index”.
If I try to change the type with “out = torch.tensor(out, dtype=torch.float32, device=self.config.device)” I only get the error message “ValueError: only one element tensors can be converted to Python scalars”.

_bertify Method

def _bertify(self, doc: Doc) -> torch.Tensor:
    subwords_batches = bert.get_subwords_batches(doc, self.config,
                                                    self.tokenizer)

    special_tokens = np.array([self.tokenizer.cls_token_id,
                                self.tokenizer.sep_token_id,
                                self.tokenizer.pad_token_id])
    subword_mask = ~(np.isin(subwords_batches, special_tokens))

    subwords_batches_tensor = torch.tensor(subwords_batches,
                                            device=self.config.device,
                                            dtype=torch.long)
    subword_mask_tensor = torch.tensor(subword_mask,
                                        device=self.config.device)

    # Obtain bert output for selected batches only
    attention_mask = (subwords_batches != self.tokenizer.pad_token_id)
    out = ""
    if "bert" in self.config.bert_model:
        out, _ = self.bert(
            subwords_batches_tensor,
            attention_mask=torch.tensor(
                attention_mask, device=self.config.device))
        del _
    elif "electra" in self.config.bert_model:
        out = self.bert(
            subwords_batches_tensor,
            attention_mask=torch.tensor(
                attention_mask, device=self.config.device))
        out = torch.tensor(out, dtype=torch.float32, device=self.config.device)
    
    # [n_subwords, bert_emb]
    return out[subword_mask_tensor]

How do I fix this?

Thanks

I’m not sure if the type really changes.
Assuming the "bert" config is the default one while "electra" is the new one, note that the output of self.bert returns a tuple while only the first element is used:

    if "bert" in self.config.bert_model:
        out, _ = self.bert(
            subwords_batches_tensor,
            attention_mask=torch.tensor(
                attention_mask, device=self.config.device))
        del _

Maybe you want to use the same approach in your "electra" branch.

Thank you for your reply. Through this I have found the solution.

bert returns a tuple with two values (out, _) and electra returns a tuple with one value (out).
So with electra the tuple is not resolved to the variables.

Solution: Adress the first value of the tuple. out[0]

elif "electra" in self.config.bert_model:
    out = self.bert(
        subwords_batches_tensor,
        attention_mask=torch.tensor(
            attention_mask, device=self.config.device))[0]