def METHOD_NAME(self, states):
    # Predict target token from decoder states
    decoder_states = states.decoder_out
    lprobs = self.model.get_normalized_probs(
        [decoder_states[:, -1:]], log_probs=True
    )
    index = lprobs.argmax(dim=-1)[0, 0].item()
    if index != self.dict['tgt'].eos_index:
        token = self.dict['tgt'].string([index])
    else:
        token = self.dict['tgt'].eos_word
    return token