if save_attention_image is not None:
input = self.test_input[:1].clone()
last = (input != self.t_nul).max(0).values.nonzero().max() + 3
if save_attention_image is not None:
input = self.test_input[:1].clone()
last = (input != self.t_nul).max(0).values.nonzero().max() + 3