UserWarning: RNN module weights are not part of single contiguous chunk of memory

[erikhuck]

Has anyone else gotten the following error when using an LSTM jumping knowledge with the GAT model: "UserWarning: RNN module weights are not part of single contiguous chunk of memory"? How do we correctly use the GAT with LSTM JK to avoid this warning? Or is this a bug?
Here is the full error message:

/mlab/data/edhu227/.miniconda3/envs/acmpp/lib/python3.10/site-packages/torch/nn/modules/rnn.py:812: UserWarning: RNN module weights are not part of single contiguous chunk of memory. This means they need to be compacted at every call, possibly greatly increasing memory usage. To compact weights again call flatten_parameters(). (Triggered internally at /opt/conda/conda-bld/pytorch_1682343995026/work/aten/src/ATen/native/cudnn/RNN.cpp:982.)
  result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,

It appears that this can be fixed by calling flatten_parameters() on the internal LSTM model. But how can we do this when we don't have access to the underlying LSTM module?

My model:

class GraphAttentionNetwork(torch.nn.Module):
    def __init__(
            self, device, n_atom_features: int, n_bond_features: int | None, gat_hidden_channels: int, gat_n_layers,
            gat_out_channels: int, gat_act_func: str, gat_norm: str, aggregation: str, gat_dropout: float, n_heads: int,
            leaky_relu_slope: float, mlp_hidden_channels: int, mlp_n_layers: int, mlp_act_func: str, mlp_norm: str, mlp_dropout: float):
        super(GraphAttentionNetwork, self).__init__()
        self.device = device
        self.gat_act_func = tgres.activation_resolver(gat_act_func)
        self.gat = tgnn.GAT(
            in_channels=n_atom_features, hidden_channels=gat_hidden_channels*n_heads, num_layers=gat_n_layers, out_channels=gat_out_channels,
            dropout=gat_dropout, v2=True, act=self.gat_act_func, act_first=False, norm=gat_norm, jk='lstm', heads=n_heads,
            edge_dim=n_bond_features, negative_slope=leaky_relu_slope, add_self_loops=False)
        if aggregation == 'max' or aggregation == 'add' or aggregation == 'mean':
            mlp_in_channels = gat_out_channels
        elif aggregation == 'max-add' or aggregation == 'max-mean' or aggregation == 'add-mean':
            mlp_in_channels = gat_out_channels * 2
        elif aggregation == 'max-add-mean':
            mlp_in_channels = gat_out_channels * 3
        else:
            raise ValueError(
                f'Invalid value for reduce_operation: {aggregation}. Valid values are max, add, mean, max-add, max-mean, add-mean,'
                f' or max-add-mean')
        self.aggregation = aggregation
        self.gat_norm = tgres.normalization_resolver(gat_norm, in_channels=mlp_in_channels)
        self.mlp = tgnn.MLP(
            in_channels=mlp_in_channels, hidden_channels=mlp_hidden_channels, num_layers=mlp_n_layers, out_channels=1, dropout=mlp_dropout,
            act=mlp_act_func, act_first=False, norm=mlp_norm)
        self.to(device)

    def forward(self, data):
        atom_features, bond_indices, bond_features, batch = data.x, data.edge_index, data.edge_attr, data.batch
        x = self.gat(x=atom_features, edge_index=bond_indices, edge_attr=bond_features)
        if self.aggregation == 'max':
            x = tgnn.global_max_pool(x, batch)
        elif self.aggregation == 'add':
            x = tgnn.global_add_pool(x, batch)
        elif self.aggregation == 'mean':
            x = tgnn.global_mean_pool(x, batch)
        elif self.aggregation == 'max-add':
            x = torch.cat([tgnn.global_max_pool(x, batch), tgnn.global_add_pool(x, batch)])
        elif self.aggregation == 'max-mean':
            x = torch.cat([tgnn.global_max_pool(x, batch), tgnn.global_mean_pool(x, batch)])
        elif self.aggregation == 'add-mean':
            x = torch.cat([tgnn.global_add_pool(x, batch), tgnn.global_mean_pool(x, batch)])
        elif self.aggregation == 'max-add-mean':
            x = torch.cat([tgnn.global_max_pool(x, batch), tgnn.global_add_pool(x, batch), tgnn.global_mean_pool(x, batch)], dim=1)
        else:
            raise ValueError(f'Invalid value for self.reduce_operation: {self.aggregation}.')
        x = self.gat_norm(x)
        x = self.gat_act_func(x)
        x = self.mlp(x)
        return torch.sigmoid(x).flatten()

Prediction code (ran after training for a number of epochs):

def predicting(model, device, loader):
    model.eval()
    total_preds = torch.Tensor()
    total_labels = torch.Tensor()
    with torch.no_grad():
        for data in loader:
            data = data.to(device)
            output = model(data)

Environment

• PyG version: 2.3.1
• PyTorch version: 2.0.1
• OS: Linux
• Python version: 3.10
• CUDA/cuDNN version: 11.8
• How you installed PyTorch and PyG: conda
• Any other relevant information: torch-scatter version: 2.1.1+pt20cu118

The code for instantiating the model:

        model = GraphAttentionNetwork(
            device=device, n_atom_features=dataset.n_atom_features, n_bond_features=dataset.n_bond_features, gat_hidden_channels=64,
            gat_n_layers=4, gat_out_channels=256, gat_act_func='relu', gat_norm='batch_norm', aggregation='max-add-mean',
            gat_dropout=0.1, n_heads=10, leaky_relu_slope=0.2, mlp_hidden_channels=128, mlp_n_layers=3, mlp_act_func='relu',
            mlp_norm='batch_norm', mlp_dropout=0.1)

[rusty1s]

Duplicate of #7999. I will try to reproduce.