Softmax separately on each graph in a mini-batch

[matanat]

I tried to look for an answer here and also in the documentation, sorry if I missed it.
I have a dataset comprising of Heterogeneous graphs, each with a different number of nodes.
I want to run Softmax on the node features of each graph separately, meaning I'm looking for a distribution over each graph's nodes.

At the moment I'm doing:

train_loader = DataLoader(train_set, ..)  

for batch in train_loader: 
   pred = model(batch.x_dict, batch.edge_index_dict) 
   
   labels = batch.y_dict['some_node'] 
   index = list()
   for i, data in enumerate(batch.to_data_list()):
       index.extend([i] * data["some_node"].num_nodes)
   index = torch.tensor(index)

   res = softmax(src=pred, index=index)

Is there a more elegant way to get the batch assignment vector for each graph in the batch?
Running batch.batch returns:
AttributeError: 'HeteroDataBatch' has no attribute 'batch'

Thanks!

[rusty1s]

There exists a batch attribute for each node type:

res = softmax(pred['some_node'], batch['some_node'].batch)

[matanat]

Thank you!