robinzixuan commited on
Commit
0608ae2
1 Parent(s): fb42e34

Update modeling_bert.py

Browse files
Files changed (1) hide show
  1. modeling_bert.py +1 -1
modeling_bert.py CHANGED
@@ -57,7 +57,7 @@ from ...utils import (
57
  from .configuration_bert import BertConfig
58
 
59
  BertModel.register_for_auto_class("AutoModel")
60
-
61
  def softmax_n_shifted_zeros(input: torch.Tensor, n: int, dim=-1) -> torch.Tensor:
62
  """
63
  $\text(softmax)_n(x_i) = exp(x_i) / (n + \sum_j exp(x_j))$
 
57
  from .configuration_bert import BertConfig
58
 
59
  BertModel.register_for_auto_class("AutoModel")
60
+ BertForMaskedLM.register_for_auto_class("AutoModelForMaskedLM")
61
  def softmax_n_shifted_zeros(input: torch.Tensor, n: int, dim=-1) -> torch.Tensor:
62
  """
63
  $\text(softmax)_n(x_i) = exp(x_i) / (n + \sum_j exp(x_j))$