torch.jit.Attribute don't keep attribute

username: Thytu
cooked: <p>Hi,</p>
<p>I have an attribute I would like to keep in my class (see below).<br>
To keep it after scripting my model (<code>module = jit.script(module)</code>), I wrapped it with <code>torch.jit.Attribute</code> but if I do that, torch can’t find this attribute while scripting it (see error below).</p>
<p><strong>Code</strong></p>
<pre><code class="lang-python">class LSTM(nn.Module):
    def __init__(self, input_size, output_size, **kwargs):
        super(LSTM, self).__init__()

        logger.debug("Initializing LSTM")

        bidirectional = kwargs.get('bidirectional', False)
        self.__num_internal_layers = torch.jit.Attribute(2 if bidirectional else 1, int) # # I defined it here

        # [...]

        self.l_lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=self.lstm_hidden_size,
            num_layers=self.lstm_num_layers,
            batch_first=True,
            bidirectional=bidirectional,
            device=self.device
        )

         # [...]

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        """
        Forward pass

        params:
            x:      tensor of shape (B, N, X) where B is the batch_size, N the sequence and X the data itself

        returns:
            y:      tensor of shape (B, N, Y) where Y is the output_size
        """

        h_0 = torch.zeros(self.lstm_num_layers * self.__num_internal_layers.value, x.size(0), self.lstm_hidden_size, device=self.device) # I use it here
        c_0 = torch.zeros(self.lstm_num_layers * self.__num_internal_layers.value, x.size(0), self.lstm_hidden_size, device=self.device) # I use it here

        # [...]

</code></pre>
<p><strong>Error</strong></p>
<pre><code class="lang-auto">Module 'LSTM' has no attribute '__num_internal_layers' :
  File "/home/thytu/Prog/Blackfoot/herding-cats-AI/src/LSTM.py", line 77
        """
    
        h_0 = torch.zeros(self.lstm_num_layers * self.__num_internal_layers.value, x.size(0), self.lstm_hidden_size, device=self.device)
                                                 ~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
        c_0 = torch.zeros(self.lstm_num_layers * self.__num_internal_layers.value, x.size(0), self.lstm_hidden_size, device=self.device)
</code></pre>
<p>Do you have any idea why?</p>

username: ptrblck
cooked: <p>Double underscores have a special meaning in Python and are used as private variables.<br>
They are not accessible via their attribute name, as the <code>_CLASSNAME</code> will be appended to them as seen here:</p>
<pre><code class="lang-python">class MyClass(object):
    def __init__(self):
        self.my_attr = 1
        self.__my_private_attr = 2
        
m = MyClass()
print(m.my_attr)
# &gt; 1
print(m.__my_private_attr)
# &gt; AttributeError: 'MyClass' object has no attribute '__my_private_attr'
print(m._MyClass__my_private_attr)
# &gt; 2
</code></pre>
<p>I’m not familiar with your use case, but you could remove the double underscores and also the <code>.value</code> call:</p>
<pre><code class="lang-python">class LSTM(nn.Module):
    def __init__(self, input_size, output_size, **kwargs):
        super(LSTM, self).__init__()
        bidirectional = kwargs.get('bidirectional', False)
        self.num_internal_layers = torch.jit.Attribute(2 if bidirectional else 1, int) # # I defined it here

        self.l_lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=1,
            num_layers=1,
            batch_first=True,
            bidirectional=bidirectional,
            device='cpu'
        )
        
    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        h_0 = torch.zeros(1 * self.num_internal_layers, x.size(0), 1, device='cpu') # I use it here
        c_0 = torch.zeros(1 * self.num_internal_layers, x.size(0), 1, device='cpu') # I use it here
        return h_0
        
model = torch.jit.script(LSTM(1, 1))
</code></pre>