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

最新推荐文章于 2024-04-02 10:06:57 发布

sunny0722

最新推荐文章于 2024-04-02 10:06:57 发布

阅读量1.1k

点赞数 1

本文链接：https://blog.csdn.net/sunny0722/article/details/106888934

版权

最近几天的代码用到了gru，但是当代码使用DataParallel时，就会报下边图片的问题。使用单GPU训练时就不存在下述警告。

查阅资料以后，发现作者https://blog.csdn.net/feifei3211/article/details/102998288的解决方法非常有效地解决了该警告。为了便于查阅，我还是将问题及问题的解决方法记录下来。解决方法如下：

class en_GRU(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers):
        super(en_GRU, self).__init__()
        self.gru = nn.GRU(input_size, hidden_size, num_layers)

    def forward(self, x, hidden):
        if not hasattr(self, '_flattened'):
            self.gru.flatten_parameters()
        setattr(self, '_flattened', True)
        output, h_n = self.gru(x, hidden)
        return output, h_n

在gru的forward函数里加入了标红的语句，self.gru.flatten_parameters()。