I love this library! Trying to perform node classification over a graph with weighted edges. I am using stochastic training and therefore blocks. The data loader is `dgl.dataloading.DataLoader`

with a `dgl.dataloading.NeighborSampler([4, 4])`

sampler. I am trying to pass the edge weights to my model, and I retrieve them as follows during a training loop:

```
x = blocks[0].srcdata['feature']
weights = blocks[0].edata['weight']
y = blocks[-1].dstdata['label']
```

With the following forward pass:

```
def forward(self, blocks, x, weights):
h = x
for l, (layer, block) in enumerate(zip(self.layers, blocks)):
print(f"Length of weights: {len(weights)}, Number of Edges: {block.number_of_edges()}")
h = layer(block, h, weights)
if l != len(self.layers) - 1:
h = F.relu(h)
h = self.dropout(h)
return h
```

Where each layer is just `dglnn.SAGEConv(hid_size, hid_size, 'mean')`

. A single forward pass for one epoch works with the edge tensor the exact same shape as the number of edges, but a second epoch returns the following error. Am I implementing the sampling incorrectly?

```
File "/Users/ajmal.aziz/opt/miniconda3/envs/kbricks/lib/python3.10/site-packages/dgl/nn/pytorch/conv/sageconv.py", line 218, in forward
assert edge_weight.shape[0] == graph.number_of_edges()
AssertionError
Length of weights: 725, Number of Edges: 725
Length of weights: 725, Number of Edges: 185
```

Any help is really appreciated!