pytorchでcannot assign module before Module.init() callが出た時

すごく簡単な話なのですが、いつかやらかして長時間溶かしそうなので備忘録です。

該当コードとエラー内容

import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, p_n_features_num, timesteps) -> None:
        self.p_n_features_num = p_n_features_num
        self.linear = nn.Linear(
            p_n_features_num, p_n_features_num, bias=True)
        self.leakyrelu = nn.LeakyReLU()

    def forward(self, input_net):
        input_net = input_net.view(input_net.size(0), self.p_n_features_num)
        return self.leakyrelu(self.linear(input_net))

ここでEncoderの__init__を呼び出すと
cannot assign module before Module.init() call
のエラーが出ます。

解決策、修正コード

これはinitのときにnn.Moduleを継承しているのでsuperメソッドを始めに呼び出さなかったのが原因です。そのため、以下のように修正してあげれば直ります。

import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, p_n_features_num, timesteps) -> None:
        super(Encoder, self).__init__()
        self.p_n_features_num = p_n_features_num
        self.linear = nn.Linear(
            p_n_features_num, p_n_features_num, bias=True)
        self.leakyrelu = nn.LeakyReLU()

    def forward(self, input_net):
        input_net = input_net.view(input_net.size(0), self.p_n_features_num)
        return self.leakyrelu(self.linear(input_net))