lhallee commited on
Commit
8f900cf
·
verified ·
1 Parent(s): c2333a7

Upload modeling_fastesm.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. modeling_fastesm.py +5 -0
modeling_fastesm.py CHANGED
@@ -377,6 +377,11 @@ class Pooler:
377
  attention_mask: Optional[torch.Tensor] = None,
378
  attentions: Optional[torch.Tensor] = None
379
  ) -> torch.Tensor:
 
 
 
 
 
380
  final_emb: List[torch.Tensor] = []
381
  for pooling_type in self.pooling_types:
382
  final_emb.append(self.pooling_options[pooling_type](emb=emb, attention_mask=attention_mask, attentions=attentions))
 
377
  attention_mask: Optional[torch.Tensor] = None,
378
  attentions: Optional[torch.Tensor] = None
379
  ) -> torch.Tensor:
380
+ if attention_mask is not None:
381
+ assert attention_mask.sum(dim=-1).min() > 0, (
382
+ "Pooler received samples with all-zero attention masks. "
383
+ "This causes NaN from division by zero. Filter empty inputs before pooling."
384
+ )
385
  final_emb: List[torch.Tensor] = []
386
  for pooling_type in self.pooling_types:
387
  final_emb.append(self.pooling_options[pooling_type](emb=emb, attention_mask=attention_mask, attentions=attentions))