# Groupby mean according to two tensor

Suppose we have the tensor looks this follow, assume that we have 3 different class.

``````sample = torch.Tensor([
[0.1, 0.1],    #-> batch 1 / class 1
[0.2, 0.2],    #-> batch 1 / class 2
[0.4, 0.4],    #-> batch 1 / class 2
[0.5, 0.5]     #-> batch 1 / class 0
[0.7, 0.7]     #-> batch 1 / class 0
[0.3, 0.3]     #-> batch 2 / class 1
[0.1, 0.1]     #-> batch 2 / class 1
[0.8, 0.8]     #-> batch 2 / class 0
])

batch_num = torch.Tensor([5, 3], dtype = torch.Long)
node_type = torch.Tensor([1, 2, 2, 0, 0, 1, 1, 0], dtype = torch.Long)
``````

`sample` is our data.

`batch_num `means how many nodes in each sample in the batch.

`node_type ` means type of each node, `len(node_type) == sample .size(0)`

What I want to compute is that:

for each sample, I want to compute the mean according to the class, so the final result would looks like

``````output= torch.Tensor([
[0.1, 0.1],    #-> batch 1 / class 1
[0.3, 0.3],    #-> batch 1 / class 2
[0.6, 0.6]     #-> batch 1 / class 0
[0.2, 0.2]     #-> batch 2 / class 1
[0.8, 0.8]     #-> batch 2 / class 0
[0., 0.]       #-> batch 2 / class 2
])
``````

Note that, for batch 2, there only two classes, for this case, we need to add `0` tensor to it

I was stuck here for a while, any help would be appreciated.