Preserve runtime errors in MessagePassing (#6417)

test/nn/conv/test_message_passing.py

@@ -100,11 +100,11 @@ def test_my_conv_out_of_bounds():
 
     conv = MyConv(8, 32)
 
-    with pytest.raises(ValueError, match="valid indices"):
+    with pytest.raises(IndexError, match="valid indices"):
         edge_index = torch.tensor([[-1, 1, 2, 2], [0, 0, 1, 1]])
         conv(x, edge_index, value)
 
-    with pytest.raises(ValueError, match="valid indices"):
+    with pytest.raises(IndexError, match="valid indices"):
         edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
         conv(x, edge_index, value)
 

torch_geometric/nn/conv/message_passing.py

@@ -238,12 +238,15 @@ def __lift__(self, src, edge_index, dim):
                 index = edge_index[dim]
                 return src.index_select(self.node_dim, index)
             except (IndexError, RuntimeError) as e:
-                if 'CUDA' in str(e):
-                    raise ValueError(
-                        f"Encountered a CUDA error. Please ensure that all "
-                        f"indices in 'edge_index' point to valid indices "
-                        f"in the interval [0, {src.size(self.node_dim)}) in "
-                        f"your node feature matrix and try again.")
+                if index.min() < 0 or index.max() >= src.size(self.node_dim):
+                    raise IndexError(
+                        f"Encountered an index error. Please ensure that all "
+                        f"indices in 'edge_index' point to valid indices in "
+                        f"the interval [0, {src.size(self.node_dim) - 1}] "
+                        f"(got interval "
+                        f"[{int(index.min())}, {int(index.max())}])")
+                else:
+                    raise e
 
                 if index.numel() > 0 and index.min() < 0:
                     raise ValueError(