Probleme whith torch_geometric.nn.TopKPooling

please I wanted to use the maximum using TopKPooling from torch_geometric.nn in this way:
in the init() part

self.pool1 = TopKPooling(8, ratio=1)

in the forward() part

x, edge_index, _, batch, _, _ = self.pool1(x, edges, None, batch)

whit batch = 256
but i got this error message

RuntimeError: The expanded size of the tensor (256) must match the existing size (1058) at non-singleton dimension 0.  Target sizes: [256].  Tensor sizes: [1058]

can you help me to solve this problem?
and thank you