fix recompute

EleutherAI · May 6, 2021 · a1ce559 · a1ce559
1 parent df6cf10
commit a1ce559
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/megatron/text_generation_utils.py b/megatron/text_generation_utils.py
@@ -226,6 +226,7 @@ def stream_tokens(neox_args, model, context_tokens: List[List[int]], eos_token_i
  # when recomputing or at first iteration all tokens are forwarded
  tokens_to_use = context_tokens[:, :token_index_to_generate]
  positions_to_use = position_ids[:, :token_index_to_generate]
+ layer_past = torch.Tensor().cuda()
  else:
  # otherwise only the last tokens are forwarded and layer past is used for other tokens
  tokens_to_use = context_tokens[:, token_index_to_generate - 1].view(batch_size, -1) # view applied to keep dimensions