Pytorch_geometric: EdgeConv implementation error in test(loader) method

[Raverss]

🐛 Bug

EdgeConv implementation line throws IndexError: The shape of the mask [50] at index 0 does not match the shape of the indexed tensor [4] at index 0.

To Reproduce

Steps to reproduce the behavior:

1. Run file.
2. Wait until the end of the first epoch for test(loader) call.

Expected behavior

Calculation of the accuracy and IoU at the given epoch.

Environment

• OS: Ubuntu 18.04
• Python version: 3.7
• PyTorch version: 1.5.0
• CUDA/cuDNN version: 10.2/7.6
• GCC version: 7.5.0

Additional context

I'm new to the PyTorch but it seems that replacing following lines:

i = intersection[j, loader.dataset.y_mask[category[j]]]
u = union[j, loader.dataset.y_mask[category[j]]]

with

i = intersection[j, loader.dataset.y_mask[category[j], :4]]
u = union[j, loader.dataset.y_mask[category[j], :4]]

fixes crashing and give reasonable results (I'm not sure how I'm new to PyTorch). I reckon that 'Airplane': [0,1,2,3] corresponds to :4. I don't know how to generalize to make it work for list of categories.

[Raverss]

Update: Proposed fix in "Additional context" doesn't work, it was just pure luck.

[rusty1s]

I will look into it.

[rusty1s]

This is now fixed in master. Sorry for the inconveniences!