Songwei Ge commited on
Commit
b5268ad
1 Parent(s): f6e53b7
Files changed (1) hide show
  1. utils/attention_utils.py +2 -2
utils/attention_utils.py CHANGED
@@ -69,8 +69,8 @@ def plot_attention_maps(atten_map_list, obj_tokens, save_dir, seed, tokens_vis=N
69
  else:
70
  axs_xlabel = ''
71
  for token_id in obj_tokens[tid]:
72
- axs_xlabel += tokens_vis[token_id.item() -
73
- 1][:-len('</w>')]
74
  axs[tid].set_title(axs_xlabel)
75
 
76
  norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
 
69
  else:
70
  axs_xlabel = ''
71
  for token_id in obj_tokens[tid]:
72
+ axs_xlabel += ' ' + tokens_vis[token_id.item() -
73
+ 1][:-len('</w>')]
74
  axs[tid].set_title(axs_xlabel)
75
 
76
  norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)