to_hetero fails at dropout due to torch.rand not accepting proxies.

[szemyd]

🐛 Describe the bug

I'm trying to pass these layers through to_hertero()

class GNNEncoder(t.nn.Module):
    def __init__(
        self,
        layers: ModuleList,
        p_dropout_edges: Optional[float],
        p_dropout_features: Optional[float],
    ):
        super().__init__()
        self.layers = layers
        self.p_dropout_edges = p_dropout_edges
        self.p_dropout_features = p_dropout_features

    def forward(self, x, edge_index):
        for index, layer in enumerate(self.layers):
            if index == len(self.layers) - 1:
                x = layer(x, edge_index)
            else:
                if self.p_dropout_edges is not None:
                    edge_index, _ = dropout_adj(
                        edge_index,
                        p=self.p_dropout_edges,
                        force_undirected=True,
                        training=self.training,
                    )
                if self.p_dropout_features is not None:
                    x = F.dropout(x, p=self.p_dropout_features, training=self.training)

                x = layer(x, edge_index).relu()

        return x

However the tracing fails at:

dropout_adj(
                        edge_index,
                        p=self.p_dropout_edges,
                        force_undirected=True,
                        training=self.training,
                    )

I'm getting a type error:
rand() received an invalid combination of arguments - got (Proxy, device=Attribute) ...

It fails at mask = torch.rand(row.size(0), device=edge_index.device) >= p

Environment

• PyG version: 2.0.5
• PyTorch version: 1.10.2
• OS: MacOS 12.3.1
• Python version: 3.9.10
• CUDA/cuDNN version:
• How you installed PyTorch and PyG (conda, pip, source): from https://github.com/pyg-team/pytorch_geometric/archive/master.zip
• Any other relevant information (e.g., version of torch-scatter):

[rusty1s]

I am afraid dropout_adj does not work yet in combination with to_hetero. Will try to look into it.

[rusty1s]

I finally had time to look into this. It looks like to_hetero() already works as expected here. The trick is that one needs to register any functions before-hand that shall not be traced by torch.fx:

torch.fx.wrap('dropout_adj')

See #4653 for an example.