'How to pop elements from a tensor in Pytorch?

I want to drop/pop elements from a tensor in Pytorch, something similar to pop operation in python. In the following code , if the condition is met, it removes two elements from the array, current and the next. I have a corresponding pytorch tensor. That means if the length of array is 10, I have the tensor last_hidden_state of size (1,10,768). After taking the mean of elements last_hidden_state[:,index-1,:], last_hidden_state[:,index,:] and last_hidden_state[:,index+1,:] I want to drop last_hidden_state[:,index,:] and last_hidden_state[:,index+1,:] Just like popping the current and next element from array. That means I should get a tensor of size (1,8,768) but with this code sometimes it returns (1,7,768) or (1,6,768). What is it that I am doing wrong? I am new to Pytorch tensors, thank you

def function_merge (prev_el, curr_el, next_el,index, array):
    if(curr_el.startswith('##') and next_el.startswith('##')):
         array[index-1] =  prev_el + curr_el + next_el
         array.pop(index) #remove current element
         array.pop(index) #remove next element

         last_hidden_state[:,index-1,:] = torch.add(last_hidden_state[:,index-1,:],last_hidden_state[:,index,:])
         last_hidden_state[:,index-1,:] = torch.add(last_hidden_state[:,index-1,:],last_hidden_state[:,index+1,:])
         last_hidden_state[:,index-1,:] = torch.mean(last_hidden_state[:,index-1,:])
            
         last_hidden_state = torch.cat((last_hidden_state[:,:index,:],last_hidden_state[:,index+2:,:]), axis=1)
         return array, last_hidden_state
               


Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source