problem with custom collate_fn in Dataloader - machine-learning

Question

I am trying to do some graph classification task. to do so I needed to write my own collate function to use in Dataloader in pytorch.
train_loader =loader.DataLoader(train_dataset, batch_size=2, shuffle=False,collate_fn=collate_mydataset)
"collate_mydataset" function seems to work well when is called directly when I pass two different samples (with different number of nodes/edges/fearures/etc ) in a list to it. it pads correctly everything. but when I set batch_size>1 ( or to make it similar to the case of calling it directly , lets say batch_size=2 ) it doesn't work well and says stack expects each tensor to be equal size, but got [32, 9] at entry 0 and [15, 9] at entry 1 ( what I understand it says node_features are not padded correctly. here is also my collate function:
def collate_mydataset(samples):
num_nodes_list = [data[5] for data in samples]
max_num_nodes = max(num_nodes_list)
num_edges_list = [data[2].size(0) for data in samples]
max_num_edges = max(num_edges_list)
features_list = [data[0] for data in samples] #node features
edge_indices_list = [data[1] for data in samples]
edge_features_list = [data[2] for data in samples]
graph_labels_list = [data[3] for data in samples]
m_list = [data[4] for data in samples]
features_padded = []
for feature in features_list:
num_nodes = feature.shape[0]
if num_nodes < max_num_nodes:
padding = torch.zeros((max_num_nodes - num_nodes, feature.shape[1]))
features_padded.append(torch.cat([feature, padding], 0))
else:
features_padded.append(feature)
features = torch.stack(features_padded, dim=0)
edge_indices_padded = []
for edge_indices in edge_indices_list:
num_edges = edge_indices.shape[1]
if num_edges < max_num_edges:
padding = torch.zeros((2, max_num_edges - num_edges))
edge_indices_padded.append(torch.cat([edge_indices, padding], 1))
else:
edge_indices_padded.append(edge_indices)
edge_indices = torch.stack(edge_indices_padded, dim=1)
edge_features_padded = []
for feature in features_list:
num_edges = feature.shape[0]
if num_edges < max_num_edges:
padding = torch.zeros((max_num_edges - num_edges, feature.shape[1]))
edge_features_padded.append(torch.cat([feature, padding], 0))
else:
edge_features_padded.append(feature)
edge_features = torch.stack(edge_features_padded, dim=0)
graph_labels = torch.stack(graph_labels_list, dim=0)
m_padded = []
for m in m_list:
num_nodes = m.shape[0]
if num_nodes < max_num_nodes:
padding = torch.zeros((max_num_nodes - num_nodes, m.shape[1]))
m_padded.append(torch.cat([m, padding], 0))
else:
m_padded.append(m)
m = torch.stack(m_padded, dim=0)
return features, edge_indices, edge_features, graph_labels, m
what could be the issue ?