derp

huggingface · Mar 14, 2024 · 10360b3 · 10360b3
1 parent 8c29e49
commit 10360b3
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
@@ -1932,7 +1932,9 @@ def _contrastive_search(
 
  # keep track of which sequences are already finished
  batch_size, cur_len = (
- model_kwargs["attention_mask"].shape if model_kwargs.get("attention_mask", None) else input_ids.shape
+ model_kwargs["attention_mask"].shape
+ if model_kwargs.get("attention_mask", None) is not None
+ else input_ids.shape
  )
  unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
  model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
@@ -2394,7 +2396,9 @@ def _greedy_search(
 
  # keep track of which sequences are already finished
  batch_size, cur_len = (
- model_kwargs["attention_mask"].shape if model_kwargs.get("attention_mask", None) else input_ids.shape
+ model_kwargs["attention_mask"].shape
+ if model_kwargs.get("attention_mask", None) is not None
+ else input_ids.shape
  )
  unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
  model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
@@ -2696,7 +2700,9 @@ def _sample(
 
  # keep track of which sequences are already finished
  batch_size, cur_len = (
- model_kwargs["attention_mask"].shape if model_kwargs.get("attention_mask", None) else input_ids.shape
+ model_kwargs["attention_mask"].shape
+ if model_kwargs.get("attention_mask", None) is not None
+ else input_ids.shape
  )
  unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
  model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
@@ -4537,7 +4543,9 @@ def _assisted_decoding(
 
  # keep track of which sequences are already finished
  batch_size, cur_len = batch_size, cur_len = (
- model_kwargs["attention_mask"].shape if model_kwargs.get("attention_mask", None) else input_ids.shape
+ model_kwargs["attention_mask"].shape
+ if model_kwargs.get("attention_mask", None) is not None
+ else input_ids.shape
  )
  unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
  model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)