<h2 dir="auto">🐛 Bug 🐛</h2>

<p dir="auto">When loading a saved sparse tensor an error occurs: <code class="notranslate">RuntimeError: size is inconsistent with indices: for dim 0, size is 3 but found index 140115406052656</code><br>
Clearly, that is impossible (can't even fit such a large tensor on my memory)</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<ol dir="auto">
<li>Install PyTorch and torch-sparse</li>
<li>Run the below code (which produces the above error for me)</li>
</ol>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
x=torch.zeros(3,3)
x[1][1]=1
x=x.to_sparse()
torch.save(x, './scratch/test4e')
y = torch.load('./scratch/test4e')"><pre class="notranslate"><code class="notranslate">import torch
x=torch.zeros(3,3)
x[1][1]=1
x=x.to_sparse()
torch.save(x, './scratch/test4e')
y = torch.load('./scratch/test4e')
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">Loads <code class="notranslate">y</code> normally ...</p>
<h2 dir="auto">Environment</h2>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="
PyTorch version: 1.4.0
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.10.2

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: GeForce GTX 1070
Nvidia driver version: 440.59
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.18.1
[pip3] torch==1.4.0
[pip3] torchvision==0.5.0
[conda] torch                     1.4.0                     &lt;pip&gt;
[conda] torch-cluster             1.4.5                     &lt;pip&gt;
[conda] torch-geometric           1.4.1                     &lt;pip&gt;
[conda] torch-scatter             2.0.3                     &lt;pip&gt;
[conda] torch-sparse              0.5.1                     &lt;pip&gt;"><pre class="notranslate"><code class="notranslate">
PyTorch version: 1.4.0
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.10.2

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: GeForce GTX 1070
Nvidia driver version: 440.59
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.18.1
[pip3] torch==1.4.0
[pip3] torchvision==0.5.0
[conda] torch                     1.4.0                     &lt;pip&gt;
[conda] torch-cluster             1.4.5                     &lt;pip&gt;
[conda] torch-geometric           1.4.1                     &lt;pip&gt;
[conda] torch-scatter             2.0.3                     &lt;pip&gt;
[conda] torch-sparse              0.5.1                     &lt;pip&gt;
</code></pre></div>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ezyang/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ezyang">@ezyang</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gchanan/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gchanan">@gchanan</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zou3519/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zou3519">@zou3519</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/vincentqb/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/vincentqb">@vincentqb</a></p>
<p dir="auto"><a href="https://github.com/pytorch/pytorch/blob/04c0971679374148aa4105e2e998de6478eca1eb/torch/nn/modules/rnn.py#L723-L727">Current implementation</a> of LSTMCell calls internal C function which makes it infeasible to write a layer normalization of LSTM cell by wrapping it with <code class="notranslate">nn.LayerNorm</code>.</p>
<p dir="auto">It would be great to add this feature.</p>
<p dir="auto">We could add trainable parameters inside <code class="notranslate">RNNCellBase</code> as follows</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class RNNCellBase(nn.Module):
    def __init__(self, input_size, hidden_size, bias, num_chunks):
        super(RNNCellBase, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias
        self.weight_ih = Parameter(torch.Tensor(num_chunks * hidden_size, input_size))
        self.weight_hh = Parameter(torch.Tensor(num_chunks * hidden_size, hidden_size))
        self.ln_ih = nn.LayerNorm(num_chunks * hidden_size)
        self.ln_hh = nn.LayerNorm(num_chunks * hidden_size)
        self.ln_c = nn.LayerNorm(hidden_size)"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">RNNCellBase</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">input_size</span>, <span class="pl-s1">hidden_size</span>, <span class="pl-s1">bias</span>, <span class="pl-s1">num_chunks</span>):
        <span class="pl-en">super</span>(<span class="pl-v">RNNCellBase</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">input_size</span> <span class="pl-c1">=</span> <span class="pl-s1">input_size</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">hidden_size</span> <span class="pl-c1">=</span> <span class="pl-s1">hidden_size</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">bias</span> <span class="pl-c1">=</span> <span class="pl-s1">bias</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">weight_ih</span> <span class="pl-c1">=</span> <span class="pl-v">Parameter</span>(<span class="pl-s1">torch</span>.<span class="pl-v">Tensor</span>(<span class="pl-s1">num_chunks</span> <span class="pl-c1">*</span> <span class="pl-s1">hidden_size</span>, <span class="pl-s1">input_size</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">weight_hh</span> <span class="pl-c1">=</span> <span class="pl-v">Parameter</span>(<span class="pl-s1">torch</span>.<span class="pl-v">Tensor</span>(<span class="pl-s1">num_chunks</span> <span class="pl-c1">*</span> <span class="pl-s1">hidden_size</span>, <span class="pl-s1">hidden_size</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">ln_ih</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">LayerNorm</span>(<span class="pl-s1">num_chunks</span> <span class="pl-c1">*</span> <span class="pl-s1">hidden_size</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">ln_hh</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">LayerNorm</span>(<span class="pl-s1">num_chunks</span> <span class="pl-c1">*</span> <span class="pl-s1">hidden_size</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">ln_c</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">LayerNorm</span>(<span class="pl-s1">hidden_size</span>)</pre></div>
<p dir="auto">However, the it is impossible to use it for the computation</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="return _VF.lstm_cell(
            input, hx,
            self.weight_ih, self.weight_hh,
            self.bias_ih, self.bias_hh,
        )"><pre class="notranslate"><span class="pl-k">return</span> <span class="pl-s1">_VF</span>.<span class="pl-en">lstm_cell</span>(
            <span class="pl-s1">input</span>, <span class="pl-s1">hx</span>,
            <span class="pl-s1">self</span>.<span class="pl-s1">weight_ih</span>, <span class="pl-s1">self</span>.<span class="pl-s1">weight_hh</span>,
            <span class="pl-s1">self</span>.<span class="pl-s1">bias_ih</span>, <span class="pl-s1">self</span>.<span class="pl-s1">bias_hh</span>,
        )</pre></div>
<p dir="auto">as it calls C functions.</p>
<p dir="auto">And writing it purely in Python, the speed would be too slow.</p>