How to allocate the value by condition when there are multiple matched cases in python?

Question

although I was concerned about the duplicated question, I couldn't find the proper answer. If there are any similar questions, please notify me.

I just want to allocate the class value by matching the condition.
In the below code, batch.edge_(label)_index[0] and batch.edge_(label)_index[1] indicate the target node and source node, respectively.
When I have the edge_class information which indicates different types of edges, I'd like to allocate the corresponding edge_class for the node linked to the already known edges.

batch = next(iter(train_loader))
batch.edge_class = batch.edge_class[batch.input_id]
batch.edge_index_class = torch.zeros(len(batch.edge_index[0]))
batch.to(device)

print(batch.edge_label_index)
print(batch.edge_class)
print(batch.edge_index)
print(batch.edge_index_class)

#tensor([[ 0,  1,  2,  3,  4,  5,  6,  7],
#        [ 8,  9, 10, 11, 12, 13, 14, 15]], device='cuda:0')
#tensor([2, 2, 3, 3, 3, 1, 2, 2], device='cuda:0')
#EdgeIndex([[16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
#            33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49,
#            50, 17, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65,
#            66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80],
#           [ 0,  0,  1,  1,  2,  2,  3,  3,  4,  4,  5,  5,  6,  6,  7,  7,  8,
#             8,  9,  9, 10, 10, 11, 11, 12, 12, 13, 13, 14, 14, 15, 15, 16, 16,
#            17, 17, 18, 18, 22, 22, 24, 24, 26, 26, 27, 27, 28, 28, 30, 30, 32,
#            32, 33, 33, 36, 36, 38, 38, 39, 39, 43, 43, 44, 44, 47, 47]],
#          device='cuda:0', sparse_size=(81, 81), nnz=66, sort_order=col)
#tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., #0., 0., 0., 0.,
#        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., #0., 0., 0., 0.,
#        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
#       device='cuda:0')

To do this, I run the code as below:

for i in range(len(batch.edge_index_class)):
    if((batch.edge_index[1][i] in batch.edge_label_index[0]) or (batch.edge_index[1][i] in batch.edge_label_index[1])):
        batch.edge_index_class[i] = batch.edge_class[(batch.edge_index[1][i]==batch.edge_label_index[0])|(batch.edge_index[1][i]==batch.edge_label_index[1])]
    else:
        batch.edge_index_class[i] = batch.edge_index_class[(batch.edge_index[1][i]==batch.edge_index[0])]

it returned the error because of the 17 node in batch.edge_index[0][1] and batch.edge_index[0][34] (i.e., duplicated target node)

#---------------------------------------------------------------------------
#RuntimeError                              Traceback (most recent call last)
#Cell In[76], line 5
#      3     batch.edge_index_class[i] = batch.edge_class[(batch.edge_index[1][i]==batch.edge_label_index[0])|(batch.edge_index[1][i]==batch.edge_label_index[1])]
#      4 else:
#----> 5     batch.edge_index_class[i] = batch.edge_index_class[(batch.edge_index[1][i]==batch.edge_index[0])]
#
#RuntimeError: expand(torch.cuda.FloatTensor{[2]}, size=[]): the number of sizes provided (0) must be greater or equal to the number of dimensions in the tensor (1)

print(batch.edge_index_class)
#tensor([2., 2., 2., 2., 3., 3., 3., 3., 3., 3., 1., 1., 2., 2., 2., 2., 2., 2.,
#        2., 2., 3., 3., 3., 3., 3., 3., 1., 1., 2., 2., 2., 2., 0., 0., 0., 0.,
#        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
#        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], device='cuda:0')

#Expected result
#tensor([2., 2., 2., 2., 3., 3., 3., 3., 3., 3., 1., 1., 2., 2., 2., 2., 2., 2.,
#        2., 2., 3., 3., 3., 3., 3., 3., 1., 1., 2., 2., 2., 2., 2., 2., 2., 2., 
#        2., 2., 3., 3., 3., 3., 1., 1., 1., 1., 2., 2., 2., 2., 2., 2., 2., 2., 
#        3., 3., 3., 3., 3., 3., 1., 1., 2., 2., 2., 2.], device='cuda:0')

How should I change the conditions in the for loop? Thank you for reading the question.