psodre commited on
Commit
675cdc4
·
verified ·
1 Parent(s): 4467aab

Upload modeling_nemotron_h_augmented.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. modeling_nemotron_h_augmented.py +1 -0
modeling_nemotron_h_augmented.py CHANGED
@@ -350,6 +350,7 @@ class NemotronHAugmentedForCausalLM(NemotronHForCausalLM):
350
  loss=loss,
351
  logits=logits,
352
  cache_params=nemotron_h_outputs.cache_params,
 
353
  hidden_states=nemotron_h_outputs.hidden_states,
354
  attentions=nemotron_h_outputs.attentions,
355
  )
 
350
  loss=loss,
351
  logits=logits,
352
  cache_params=nemotron_h_outputs.cache_params,
353
+ past_key_values=nemotron_h_outputs.cache_params, # HF generate threads this
354
  hidden_states=nemotron_h_outputs.hidden_states,
355
  attentions=nemotron_h_outputs.attentions,
356
  )