Update README.md

README.md

@@ -79,7 +79,7 @@ def generate_text(model_path, prompt, tokenizer, max_gen_tokens, total_sequence,
     inputs_dict = {}
     inputs_dict['input_ids'] = actual_input[:, :window].reshape(1, window).numpy()
     inputs_dict['attention_mask'] = first_attention
-    inputs_dict['tree_attention'] = np.triu(-65504*np.ones(pads), k= 1).astype('float16').reshape(1, 1, total_sequence, total_sequence)
+    inputs_dict['tree_attention'] = np.triu(-65504*np.ones(total_sequence), k= 1).astype('float16').reshape(1, 1, total_sequence, total_sequence)
     for name in inputs_names:
         if name == 'input_ids' or name == 'attention_mask' or name == 'tree_attention': continue
         inputs_dict[name] = np.zeros([1, n_heads, context-window, 128], dtype="float16")