def get_pos_states(self, input_states):
    return torch.arange(input_states.size(1)).view(1, (- 1)).repeat(input_states.size(0), 1).to(device=input_states.device)
def get_pos_states(self, input_states):
    return torch.arange(input_states.size(1)).view(1, (- 1)).repeat(input_states.size(0), 1).to(device=input_states.device)