Fix mask None-ness

#5
by cyrilvallez HF Staff - opened
Files changed (1) hide show
  1. custom_generate/generate.py +2 -1
custom_generate/generate.py CHANGED
@@ -193,7 +193,8 @@ def _contrastive_search(
193
  if "decoder_attention_mask" in model_kwargs and model_kwargs["decoder_attention_mask"] is not None:
194
  cosine_matrix_mask = model_kwargs["decoder_attention_mask"]
195
  else:
196
- cosine_matrix_mask = model_kwargs["attention_mask"]
 
197
  cosine_matrix_mask = cosine_matrix_mask.repeat_interleave(top_k, dim=0)
198
 
199
  this_peer_finished = False
 
193
  if "decoder_attention_mask" in model_kwargs and model_kwargs["decoder_attention_mask"] is not None:
194
  cosine_matrix_mask = model_kwargs["decoder_attention_mask"]
195
  else:
196
+ if model_kwargs["attention_mask"] is not None:
197
+ cosine_matrix_mask = model_kwargs["attention_mask"]
198
  cosine_matrix_mask = cosine_matrix_mask.repeat_interleave(top_k, dim=0)
199
 
200
  this_peer_finished = False