<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">When I use a single GPU, all parameters in the model get non-zero gradients;<br>
but after wrapping my model with nn.DataParallel, the gradients of some part (RNN in my case) of the model become zero.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import sys

import torch
from torch import nn
import torch.nn.functional as F
from torchexp.stat import RunningAvgDict
from torch.nn.utils.rnn import pack_padded_sequence
from torch.nn.utils.rnn import pad_packed_sequence

import numpy as np
import math


gpus = list(map(int, sys.argv[1].split(',')))


def to_device(m, x):
    &quot;&quot;&quot;Send tensor into the device of the module.

    Args:
        m (torch.nn.Module): Torch module.
        x (Tensor): Torch tensor.

    Returns:
        Tensor: Torch tensor located in the same place as torch module.

    &quot;&quot;&quot;
    assert isinstance(m, torch.nn.Module)
    device = next(m.parameters()).device
    return x.to(device)


def reset_backward_rnn_state(states):
    &quot;&quot;&quot;Sets backward BRNN states to zeroes - useful in processing of sliding windows over the inputs&quot;&quot;&quot;
    if isinstance(states, (list, tuple)):
        for state in states:
            state[1::2] = 0.
    else:
        states[1::2] = 0.
    return states


class RNNP(nn.Module):
    &quot;&quot;&quot;RNN with projection layer module

    :param int idim: dimension of inputs, default: odim of CNN
    :param int nlayers: Number of layers
    :param int enc_dim: hidden dim of BLSTMP
    :param int proj_dim: proj. dim of BLSTMP
    :param int odim: ouput of encoder dimension (default: proj_dim)
    &quot;&quot;&quot;

    def __init__(self, idim, nlayers, enc_dim, proj_dim, odim):
        super(RNNP, self).__init__()
        self.enc_dim = enc_dim
        self.proj_dim = proj_dim
        self.odim = odim
        self.nlayers = nlayers

        self.rnn0 = nn.LSTM(idim, enc_dim, num_layers=1, bidirectional=True, batch_first=True)
        self.bt0 = nn.Linear(2 * enc_dim, proj_dim)


        for i in range(1,self.nlayers):
            rnn = nn.LSTM(proj_dim, enc_dim, num_layers=1, bidirectional=True, batch_first=True)
            setattr(self, f&quot;rnn{i}&quot;, rnn)

            if i == self.nlayers - 1:
                setattr(self, f&quot;bt{i}&quot;, nn.Linear(2 * enc_dim, odim))
            else:
                setattr(self, f&quot;bt{i}&quot;, nn.Linear(2 * enc_dim, proj_dim))


    def forward(self, xs_pad, enc_lens, prev_state=None):
        &quot;&quot;&quot;RNNP forward

        :param torch.Tensor xs_pad: batch of padded input sequences (B, Tmax, cnn_o_dim)
        :param list of int enc_lens
        :param torch.Tensor prev_state: batch of previous RNN states
        :return: batch of hidden state sequences (B, Tmax, odim)
        :rtype: torch.Tensor
        &quot;&quot;&quot;
        hid_states = []

        for i in range(self.nlayers):
            total_length = xs_pad.size(1)
            xs_pack = pack_padded_sequence(xs_pad, enc_lens, batch_first = True)
            rnn = getattr(self, f&quot;rnn{i}&quot;)
            rnn.flatten_parameters()
            if prev_state is not None:
                prev_state = reset_backward_rnn_state(prev_state)
            ys, states = rnn(xs_pack, hx=None if prev_state is None else prev_state[i])
            hid_states.append(states)

            ys_pad, enc_lens= pad_packed_sequence(ys, batch_first=True, total_length=total_length)
            # ys_pad: (B, T, enc_dim)

            projected = getattr(self, f&quot;bt{i}&quot;)(ys_pad.contiguous().view(-1, ys_pad.size(2))) #(B*T, proj_dim)
            xs_pad = torch.tanh(projected.view(ys_pad.size(0), ys_pad.size(1), -1))

        return xs_pad, to_device(self,enc_lens), hid_states


class BlstmEncoder(nn.Module):
    &quot;&quot;&quot;
    :param int idim: dim of input
    :param int enc_dim: hidden dim for BLSTM
    :param int proj_dim: projection dim for BLSTMP
    :param int odim: encoder output dimension (usually set as proj_dim)
    &quot;&quot;&quot;

    def __init__(self, idim, enc_dim, proj_dim, odim):
        super(BlstmEncoder, self).__init__()
        self.idim = idim
        self.enc_dim = enc_dim
        self.proj_dim = proj_dim
        self.odim = odim

        self.cnn_model = nn.Sequential(
            nn.Conv2d(1, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2, ceil_mode=True),
        )
        cnn_o_dim = np.ceil(np.array(idim, dtype=np.float32) / 2)
        cnn_o_dim = int(cnn_o_dim) * 128

        self.nlayers = 1
        self.blstm = RNNP(cnn_o_dim, self.nlayers, enc_dim, proj_dim, odim)


    def forward(self, xs_pad, ilens, prev_state=None):
        xs_pad = xs_pad.view(xs_pad.size(0), xs_pad.size(1), 1, xs_pad.size(2)).transpose(1,2)
        xs_pad = self.cnn_model(xs_pad)

        if torch.is_tensor(ilens):
            ilens = ilens.cpu().numpy()
        else:
            ilens = np.array(ilens, dtype=np.float32)

        enc_lens = np.array(np.ceil(ilens/2), dtype=np.int64)
        enc_lens = np.array(np.ceil(np.array(enc_lens, dtype=np.float32)/2), dtype=np.int64).tolist()

        xs_pad = xs_pad.transpose(1,2)
        xs_pad = xs_pad.contiguous().view(xs_pad.size(0), xs_pad.size(1), xs_pad.size(2) * xs_pad.size(3))

        out, enc_lens, cur_state = self.blstm(xs_pad, enc_lens, prev_state)

        return out, enc_lens, cur_state


class MultiHeadBLSTM(nn.Module):

    def __init__(self):
        super(MultiHeadBLSTM, self).__init__()

        self.idim = 20
        self.odim = 10

        enc_o_dim = 128

        self.encoder = BlstmEncoder(self.idim, enc_o_dim, enc_o_dim, enc_o_dim)
        self.head = nn.Linear(enc_o_dim, self.odim)


    @property
    def device(self):
        return next(self.parameters()).device


    def forward(self, xs_pad, ilens):
        assert xs_pad.size(0) == ilens.size(0), &quot;Batch size mismatch&quot;
        batch_size = xs_pad.size(0)

        # Put data to device
        xs_pad = to_device(self,xs_pad)
        ilens = to_device(self,ilens)

        enc, enc_lens, _ = self.encoder(xs_pad, ilens)
        out = self.head(enc)

        return out, enc_lens


def run_batch(asr_model, ctc_loss, x, ilens, y_true, olens):

    pred, enc_lens = asr_model(x, ilens)
    if len(gpus) == 1:
        olens = to_device(asr_model, olens)
    else:
        olens = to_device(asr_model.module, olens)
    pred = F.log_softmax(pred, dim=-1)

    loss = ctc_loss(pred.transpose(0,1).contiguous(),
                         y_true.cuda().to(dtype=torch.long),
                         enc_lens.cpu().to(dtype=torch.long),
                         olens.cpu().to(dtype=torch.long))
    print('ctc loss: ', loss)

    info = { 'loss': loss.item() }
    loss.backward()

    return info


def run_task(asr_model, ctc_loss, x, ilens, ys, olens):
    batch_size = len(ys)
    info = run_batch(asr_model, ctc_loss, x, ilens, ys, olens)

    # torch.set_printoptions(profile=&quot;full&quot;)
    if len(gpus) == 1:
        count = 0
        for n, p in asr_model.named_parameters():
            print('grad of param {}, {}: '.format(count, n), p.grad)#[0, 0, 0])
            count += 1
        print('num of params with grad: ', count)
    else:
        count = 0
        for n, p in asr_model.module.named_parameters():
            print('grad of param {}, {}: '.format(count, n), p.grad)#[0, 0, 0])
            count += 1
        print('num of params with grad: ', count)


if __name__ == '__main__':
    asr_model = MultiHeadBLSTM().cuda()
    if len(gpus) &gt; 1:
        asr_model = nn.DataParallel(asr_model, device_ids=gpus)

    ctc_loss = nn.CTCLoss(blank=0, reduction='mean', zero_infinity=True)

    x = torch.rand(4, 20, 20)#.cuda() 
    ilens = torch.tensor([20, 20, 20, 20])#.cuda()
    ys = [torch.tensor([1, 1, 1]),
          torch.tensor([1, 1, 1, 1]),
          torch.tensor([1, 1, 1, 1, 1]),
          torch.tensor([1, 1, 1, 1, 1, 1])]
    olens = torch.tensor([3, 4, 5, 6])#.cuda()

    y_true = torch.cat(ys)#.cuda()

    run_task(asr_model, ctc_loss, x, ilens, y_true, olens)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">sys</span>

<span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span> <span class="pl-k">import</span> <span class="pl-s1">nn</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-s1">functional</span> <span class="pl-k">as</span> <span class="pl-v">F</span>
<span class="pl-k">from</span> <span class="pl-s1">torchexp</span>.<span class="pl-s1">stat</span> <span class="pl-k">import</span> <span class="pl-v">RunningAvgDict</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-s1">utils</span>.<span class="pl-s1">rnn</span> <span class="pl-k">import</span> <span class="pl-s1">pack_padded_sequence</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-s1">utils</span>.<span class="pl-s1">rnn</span> <span class="pl-k">import</span> <span class="pl-s1">pad_packed_sequence</span>

<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>
<span class="pl-k">import</span> <span class="pl-s1">math</span>


<span class="pl-s1">gpus</span> <span class="pl-c1">=</span> <span class="pl-en">list</span>(<span class="pl-en">map</span>(<span class="pl-s1">int</span>, <span class="pl-s1">sys</span>.<span class="pl-s1">argv</span>[<span class="pl-c1">1</span>].<span class="pl-en">split</span>(<span class="pl-s">','</span>)))


<span class="pl-k">def</span> <span class="pl-en">to_device</span>(<span class="pl-s1">m</span>, <span class="pl-s1">x</span>):
    <span class="pl-s">"""Send tensor into the device of the module.</span>
<span class="pl-s"></span>
<span class="pl-s">    Args:</span>
<span class="pl-s">        m (torch.nn.Module): Torch module.</span>
<span class="pl-s">        x (Tensor): Torch tensor.</span>
<span class="pl-s"></span>
<span class="pl-s">    Returns:</span>
<span class="pl-s">        Tensor: Torch tensor located in the same place as torch module.</span>
<span class="pl-s"></span>
<span class="pl-s">    """</span>
    <span class="pl-k">assert</span> <span class="pl-en">isinstance</span>(<span class="pl-s1">m</span>, <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>)
    <span class="pl-s1">device</span> <span class="pl-c1">=</span> <span class="pl-en">next</span>(<span class="pl-s1">m</span>.<span class="pl-en">parameters</span>()).<span class="pl-s1">device</span>
    <span class="pl-k">return</span> <span class="pl-s1">x</span>.<span class="pl-en">to</span>(<span class="pl-s1">device</span>)


<span class="pl-k">def</span> <span class="pl-en">reset_backward_rnn_state</span>(<span class="pl-s1">states</span>):
    <span class="pl-s">"""Sets backward BRNN states to zeroes - useful in processing of sliding windows over the inputs"""</span>
    <span class="pl-k">if</span> <span class="pl-en">isinstance</span>(<span class="pl-s1">states</span>, (<span class="pl-s1">list</span>, <span class="pl-s1">tuple</span>)):
        <span class="pl-k">for</span> <span class="pl-s1">state</span> <span class="pl-c1">in</span> <span class="pl-s1">states</span>:
            <span class="pl-s1">state</span>[<span class="pl-c1">1</span>::<span class="pl-c1">2</span>] <span class="pl-c1">=</span> <span class="pl-c1">0.</span>
    <span class="pl-k">else</span>:
        <span class="pl-s1">states</span>[<span class="pl-c1">1</span>::<span class="pl-c1">2</span>] <span class="pl-c1">=</span> <span class="pl-c1">0.</span>
    <span class="pl-k">return</span> <span class="pl-s1">states</span>


<span class="pl-k">class</span> <span class="pl-v">RNNP</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>):
    <span class="pl-s">"""RNN with projection layer module</span>
<span class="pl-s"></span>
<span class="pl-s">    :param int idim: dimension of inputs, default: odim of CNN</span>
<span class="pl-s">    :param int nlayers: Number of layers</span>
<span class="pl-s">    :param int enc_dim: hidden dim of BLSTMP</span>
<span class="pl-s">    :param int proj_dim: proj. dim of BLSTMP</span>
<span class="pl-s">    :param int odim: ouput of encoder dimension (default: proj_dim)</span>
<span class="pl-s">    """</span>

    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">idim</span>, <span class="pl-s1">nlayers</span>, <span class="pl-s1">enc_dim</span>, <span class="pl-s1">proj_dim</span>, <span class="pl-s1">odim</span>):
        <span class="pl-en">super</span>(<span class="pl-v">RNNP</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">enc_dim</span> <span class="pl-c1">=</span> <span class="pl-s1">enc_dim</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">proj_dim</span> <span class="pl-c1">=</span> <span class="pl-s1">proj_dim</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">odim</span> <span class="pl-c1">=</span> <span class="pl-s1">odim</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">nlayers</span> <span class="pl-c1">=</span> <span class="pl-s1">nlayers</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">rnn0</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">LSTM</span>(<span class="pl-s1">idim</span>, <span class="pl-s1">enc_dim</span>, <span class="pl-s1">num_layers</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">bidirectional</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">bt0</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-c1">2</span> <span class="pl-c1">*</span> <span class="pl-s1">enc_dim</span>, <span class="pl-s1">proj_dim</span>)


        <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">1</span>,<span class="pl-s1">self</span>.<span class="pl-s1">nlayers</span>):
            <span class="pl-s1">rnn</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">LSTM</span>(<span class="pl-s1">proj_dim</span>, <span class="pl-s1">enc_dim</span>, <span class="pl-s1">num_layers</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">bidirectional</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
            <span class="pl-en">setattr</span>(<span class="pl-s1">self</span>, <span class="pl-s">f"rnn<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">i</span><span class="pl-kos">}</span></span>"</span>, <span class="pl-s1">rnn</span>)

            <span class="pl-k">if</span> <span class="pl-s1">i</span> <span class="pl-c1">==</span> <span class="pl-s1">self</span>.<span class="pl-s1">nlayers</span> <span class="pl-c1">-</span> <span class="pl-c1">1</span>:
                <span class="pl-en">setattr</span>(<span class="pl-s1">self</span>, <span class="pl-s">f"bt<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">i</span><span class="pl-kos">}</span></span>"</span>, <span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-c1">2</span> <span class="pl-c1">*</span> <span class="pl-s1">enc_dim</span>, <span class="pl-s1">odim</span>))
            <span class="pl-k">else</span>:
                <span class="pl-en">setattr</span>(<span class="pl-s1">self</span>, <span class="pl-s">f"bt<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">i</span><span class="pl-kos">}</span></span>"</span>, <span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-c1">2</span> <span class="pl-c1">*</span> <span class="pl-s1">enc_dim</span>, <span class="pl-s1">proj_dim</span>))


    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">xs_pad</span>, <span class="pl-s1">enc_lens</span>, <span class="pl-s1">prev_state</span><span class="pl-c1">=</span><span class="pl-c1">None</span>):
        <span class="pl-s">"""RNNP forward</span>
<span class="pl-s"></span>
<span class="pl-s">        :param torch.Tensor xs_pad: batch of padded input sequences (B, Tmax, cnn_o_dim)</span>
<span class="pl-s">        :param list of int enc_lens</span>
<span class="pl-s">        :param torch.Tensor prev_state: batch of previous RNN states</span>
<span class="pl-s">        :return: batch of hidden state sequences (B, Tmax, odim)</span>
<span class="pl-s">        :rtype: torch.Tensor</span>
<span class="pl-s">        """</span>
        <span class="pl-s1">hid_states</span> <span class="pl-c1">=</span> []

        <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">self</span>.<span class="pl-s1">nlayers</span>):
            <span class="pl-s1">total_length</span> <span class="pl-c1">=</span> <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">1</span>)
            <span class="pl-s1">xs_pack</span> <span class="pl-c1">=</span> <span class="pl-en">pack_padded_sequence</span>(<span class="pl-s1">xs_pad</span>, <span class="pl-s1">enc_lens</span>, <span class="pl-s1">batch_first</span> <span class="pl-c1">=</span> <span class="pl-c1">True</span>)
            <span class="pl-s1">rnn</span> <span class="pl-c1">=</span> <span class="pl-en">getattr</span>(<span class="pl-s1">self</span>, <span class="pl-s">f"rnn<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">i</span><span class="pl-kos">}</span></span>"</span>)
            <span class="pl-s1">rnn</span>.<span class="pl-en">flatten_parameters</span>()
            <span class="pl-k">if</span> <span class="pl-s1">prev_state</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
                <span class="pl-s1">prev_state</span> <span class="pl-c1">=</span> <span class="pl-en">reset_backward_rnn_state</span>(<span class="pl-s1">prev_state</span>)
            <span class="pl-s1">ys</span>, <span class="pl-s1">states</span> <span class="pl-c1">=</span> <span class="pl-en">rnn</span>(<span class="pl-s1">xs_pack</span>, <span class="pl-s1">hx</span><span class="pl-c1">=</span><span class="pl-c1">None</span> <span class="pl-k">if</span> <span class="pl-s1">prev_state</span> <span class="pl-c1">is</span> <span class="pl-c1">None</span> <span class="pl-k">else</span> <span class="pl-s1">prev_state</span>[<span class="pl-s1">i</span>])
            <span class="pl-s1">hid_states</span>.<span class="pl-en">append</span>(<span class="pl-s1">states</span>)

            <span class="pl-s1">ys_pad</span>, <span class="pl-s1">enc_lens</span><span class="pl-c1">=</span> <span class="pl-en">pad_packed_sequence</span>(<span class="pl-s1">ys</span>, <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">total_length</span><span class="pl-c1">=</span><span class="pl-s1">total_length</span>)
            <span class="pl-c"># ys_pad: (B, T, enc_dim)</span>

            <span class="pl-s1">projected</span> <span class="pl-c1">=</span> <span class="pl-en">getattr</span>(<span class="pl-s1">self</span>, <span class="pl-s">f"bt<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">i</span><span class="pl-kos">}</span></span>"</span>)(<span class="pl-s1">ys_pad</span>.<span class="pl-en">contiguous</span>().<span class="pl-en">view</span>(<span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">ys_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">2</span>))) <span class="pl-c">#(B*T, proj_dim)</span>
            <span class="pl-s1">xs_pad</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">tanh</span>(<span class="pl-s1">projected</span>.<span class="pl-en">view</span>(<span class="pl-s1">ys_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>), <span class="pl-s1">ys_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">1</span>), <span class="pl-c1">-</span><span class="pl-c1">1</span>))

        <span class="pl-k">return</span> <span class="pl-s1">xs_pad</span>, <span class="pl-en">to_device</span>(<span class="pl-s1">self</span>,<span class="pl-s1">enc_lens</span>), <span class="pl-s1">hid_states</span>


<span class="pl-k">class</span> <span class="pl-v">BlstmEncoder</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>):
    <span class="pl-s">"""</span>
<span class="pl-s">    :param int idim: dim of input</span>
<span class="pl-s">    :param int enc_dim: hidden dim for BLSTM</span>
<span class="pl-s">    :param int proj_dim: projection dim for BLSTMP</span>
<span class="pl-s">    :param int odim: encoder output dimension (usually set as proj_dim)</span>
<span class="pl-s">    """</span>

    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">idim</span>, <span class="pl-s1">enc_dim</span>, <span class="pl-s1">proj_dim</span>, <span class="pl-s1">odim</span>):
        <span class="pl-en">super</span>(<span class="pl-v">BlstmEncoder</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">idim</span> <span class="pl-c1">=</span> <span class="pl-s1">idim</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">enc_dim</span> <span class="pl-c1">=</span> <span class="pl-s1">enc_dim</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">proj_dim</span> <span class="pl-c1">=</span> <span class="pl-s1">proj_dim</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">odim</span> <span class="pl-c1">=</span> <span class="pl-s1">odim</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">cnn_model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Sequential</span>(
            <span class="pl-s1">nn</span>.<span class="pl-v">Conv2d</span>(<span class="pl-c1">1</span>, <span class="pl-c1">128</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span><span class="pl-c1">3</span>, <span class="pl-s1">stride</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-c1">1</span>),
            <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>(),
            <span class="pl-s1">nn</span>.<span class="pl-v">Conv2d</span>(<span class="pl-c1">128</span>, <span class="pl-c1">128</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span><span class="pl-c1">3</span>, <span class="pl-s1">stride</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-c1">1</span>),
            <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>(),
            <span class="pl-s1">nn</span>.<span class="pl-v">MaxPool2d</span>(<span class="pl-c1">2</span>, <span class="pl-s1">stride</span><span class="pl-c1">=</span><span class="pl-c1">2</span>, <span class="pl-s1">ceil_mode</span><span class="pl-c1">=</span><span class="pl-c1">True</span>),
        )
        <span class="pl-s1">cnn_o_dim</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">ceil</span>(<span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">idim</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">np</span>.<span class="pl-s1">float32</span>) <span class="pl-c1">/</span> <span class="pl-c1">2</span>)
        <span class="pl-s1">cnn_o_dim</span> <span class="pl-c1">=</span> <span class="pl-en">int</span>(<span class="pl-s1">cnn_o_dim</span>) <span class="pl-c1">*</span> <span class="pl-c1">128</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">nlayers</span> <span class="pl-c1">=</span> <span class="pl-c1">1</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">blstm</span> <span class="pl-c1">=</span> <span class="pl-v">RNNP</span>(<span class="pl-s1">cnn_o_dim</span>, <span class="pl-s1">self</span>.<span class="pl-s1">nlayers</span>, <span class="pl-s1">enc_dim</span>, <span class="pl-s1">proj_dim</span>, <span class="pl-s1">odim</span>)


    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">xs_pad</span>, <span class="pl-s1">ilens</span>, <span class="pl-s1">prev_state</span><span class="pl-c1">=</span><span class="pl-c1">None</span>):
        <span class="pl-s1">xs_pad</span> <span class="pl-c1">=</span> <span class="pl-s1">xs_pad</span>.<span class="pl-en">view</span>(<span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>), <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">1</span>), <span class="pl-c1">1</span>, <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">2</span>)).<span class="pl-en">transpose</span>(<span class="pl-c1">1</span>,<span class="pl-c1">2</span>)
        <span class="pl-s1">xs_pad</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">cnn_model</span>(<span class="pl-s1">xs_pad</span>)

        <span class="pl-k">if</span> <span class="pl-s1">torch</span>.<span class="pl-en">is_tensor</span>(<span class="pl-s1">ilens</span>):
            <span class="pl-s1">ilens</span> <span class="pl-c1">=</span> <span class="pl-s1">ilens</span>.<span class="pl-en">cpu</span>().<span class="pl-en">numpy</span>()
        <span class="pl-k">else</span>:
            <span class="pl-s1">ilens</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">ilens</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">np</span>.<span class="pl-s1">float32</span>)

        <span class="pl-s1">enc_lens</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">np</span>.<span class="pl-en">ceil</span>(<span class="pl-s1">ilens</span><span class="pl-c1">/</span><span class="pl-c1">2</span>), <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">np</span>.<span class="pl-s1">int64</span>)
        <span class="pl-s1">enc_lens</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">np</span>.<span class="pl-en">ceil</span>(<span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">enc_lens</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">np</span>.<span class="pl-s1">float32</span>)<span class="pl-c1">/</span><span class="pl-c1">2</span>), <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">np</span>.<span class="pl-s1">int64</span>).<span class="pl-en">tolist</span>()

        <span class="pl-s1">xs_pad</span> <span class="pl-c1">=</span> <span class="pl-s1">xs_pad</span>.<span class="pl-en">transpose</span>(<span class="pl-c1">1</span>,<span class="pl-c1">2</span>)
        <span class="pl-s1">xs_pad</span> <span class="pl-c1">=</span> <span class="pl-s1">xs_pad</span>.<span class="pl-en">contiguous</span>().<span class="pl-en">view</span>(<span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>), <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">1</span>), <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">2</span>) <span class="pl-c1">*</span> <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">3</span>))

        <span class="pl-s1">out</span>, <span class="pl-s1">enc_lens</span>, <span class="pl-s1">cur_state</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">blstm</span>(<span class="pl-s1">xs_pad</span>, <span class="pl-s1">enc_lens</span>, <span class="pl-s1">prev_state</span>)

        <span class="pl-k">return</span> <span class="pl-s1">out</span>, <span class="pl-s1">enc_lens</span>, <span class="pl-s1">cur_state</span>


<span class="pl-k">class</span> <span class="pl-v">MultiHeadBLSTM</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-en">super</span>(<span class="pl-v">MultiHeadBLSTM</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">idim</span> <span class="pl-c1">=</span> <span class="pl-c1">20</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">odim</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>

        <span class="pl-s1">enc_o_dim</span> <span class="pl-c1">=</span> <span class="pl-c1">128</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">encoder</span> <span class="pl-c1">=</span> <span class="pl-v">BlstmEncoder</span>(<span class="pl-s1">self</span>.<span class="pl-s1">idim</span>, <span class="pl-s1">enc_o_dim</span>, <span class="pl-s1">enc_o_dim</span>, <span class="pl-s1">enc_o_dim</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">head</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-s1">enc_o_dim</span>, <span class="pl-s1">self</span>.<span class="pl-s1">odim</span>)


    <span class="pl-en">@<span class="pl-s1">property</span></span>
    <span class="pl-k">def</span> <span class="pl-en">device</span>(<span class="pl-s1">self</span>):
        <span class="pl-k">return</span> <span class="pl-en">next</span>(<span class="pl-s1">self</span>.<span class="pl-en">parameters</span>()).<span class="pl-s1">device</span>


    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">xs_pad</span>, <span class="pl-s1">ilens</span>):
        <span class="pl-k">assert</span> <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>) <span class="pl-c1">==</span> <span class="pl-s1">ilens</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>), <span class="pl-s">"Batch size mismatch"</span>
        <span class="pl-s1">batch_size</span> <span class="pl-c1">=</span> <span class="pl-s1">xs_pad</span>.<span class="pl-en">size</span>(<span class="pl-c1">0</span>)

        <span class="pl-c"># Put data to device</span>
        <span class="pl-s1">xs_pad</span> <span class="pl-c1">=</span> <span class="pl-en">to_device</span>(<span class="pl-s1">self</span>,<span class="pl-s1">xs_pad</span>)
        <span class="pl-s1">ilens</span> <span class="pl-c1">=</span> <span class="pl-en">to_device</span>(<span class="pl-s1">self</span>,<span class="pl-s1">ilens</span>)

        <span class="pl-s1">enc</span>, <span class="pl-s1">enc_lens</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">encoder</span>(<span class="pl-s1">xs_pad</span>, <span class="pl-s1">ilens</span>)
        <span class="pl-s1">out</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">head</span>(<span class="pl-s1">enc</span>)

        <span class="pl-k">return</span> <span class="pl-s1">out</span>, <span class="pl-s1">enc_lens</span>


<span class="pl-k">def</span> <span class="pl-en">run_batch</span>(<span class="pl-s1">asr_model</span>, <span class="pl-s1">ctc_loss</span>, <span class="pl-s1">x</span>, <span class="pl-s1">ilens</span>, <span class="pl-s1">y_true</span>, <span class="pl-s1">olens</span>):

    <span class="pl-s1">pred</span>, <span class="pl-s1">enc_lens</span> <span class="pl-c1">=</span> <span class="pl-en">asr_model</span>(<span class="pl-s1">x</span>, <span class="pl-s1">ilens</span>)
    <span class="pl-k">if</span> <span class="pl-en">len</span>(<span class="pl-s1">gpus</span>) <span class="pl-c1">==</span> <span class="pl-c1">1</span>:
        <span class="pl-s1">olens</span> <span class="pl-c1">=</span> <span class="pl-en">to_device</span>(<span class="pl-s1">asr_model</span>, <span class="pl-s1">olens</span>)
    <span class="pl-k">else</span>:
        <span class="pl-s1">olens</span> <span class="pl-c1">=</span> <span class="pl-en">to_device</span>(<span class="pl-s1">asr_model</span>.<span class="pl-s1">module</span>, <span class="pl-s1">olens</span>)
    <span class="pl-s1">pred</span> <span class="pl-c1">=</span> <span class="pl-v">F</span>.<span class="pl-en">log_softmax</span>(<span class="pl-s1">pred</span>, <span class="pl-s1">dim</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)

    <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-en">ctc_loss</span>(<span class="pl-s1">pred</span>.<span class="pl-en">transpose</span>(<span class="pl-c1">0</span>,<span class="pl-c1">1</span>).<span class="pl-en">contiguous</span>(),
                         <span class="pl-s1">y_true</span>.<span class="pl-en">cuda</span>().<span class="pl-en">to</span>(<span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">torch</span>.<span class="pl-s1">long</span>),
                         <span class="pl-s1">enc_lens</span>.<span class="pl-en">cpu</span>().<span class="pl-en">to</span>(<span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">torch</span>.<span class="pl-s1">long</span>),
                         <span class="pl-s1">olens</span>.<span class="pl-en">cpu</span>().<span class="pl-en">to</span>(<span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">torch</span>.<span class="pl-s1">long</span>))
    <span class="pl-en">print</span>(<span class="pl-s">'ctc loss: '</span>, <span class="pl-s1">loss</span>)

    <span class="pl-s1">info</span> <span class="pl-c1">=</span> { <span class="pl-s">'loss'</span>: <span class="pl-s1">loss</span>.<span class="pl-en">item</span>() }
    <span class="pl-s1">loss</span>.<span class="pl-en">backward</span>()

    <span class="pl-k">return</span> <span class="pl-s1">info</span>


<span class="pl-k">def</span> <span class="pl-en">run_task</span>(<span class="pl-s1">asr_model</span>, <span class="pl-s1">ctc_loss</span>, <span class="pl-s1">x</span>, <span class="pl-s1">ilens</span>, <span class="pl-s1">ys</span>, <span class="pl-s1">olens</span>):
    <span class="pl-s1">batch_size</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>(<span class="pl-s1">ys</span>)
    <span class="pl-s1">info</span> <span class="pl-c1">=</span> <span class="pl-en">run_batch</span>(<span class="pl-s1">asr_model</span>, <span class="pl-s1">ctc_loss</span>, <span class="pl-s1">x</span>, <span class="pl-s1">ilens</span>, <span class="pl-s1">ys</span>, <span class="pl-s1">olens</span>)

    <span class="pl-c"># torch.set_printoptions(profile="full")</span>
    <span class="pl-k">if</span> <span class="pl-en">len</span>(<span class="pl-s1">gpus</span>) <span class="pl-c1">==</span> <span class="pl-c1">1</span>:
        <span class="pl-s1">count</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>
        <span class="pl-k">for</span> <span class="pl-s1">n</span>, <span class="pl-s1">p</span> <span class="pl-c1">in</span> <span class="pl-s1">asr_model</span>.<span class="pl-en">named_parameters</span>():
            <span class="pl-en">print</span>(<span class="pl-s">'grad of param {}, {}: '</span>.<span class="pl-en">format</span>(<span class="pl-s1">count</span>, <span class="pl-s1">n</span>), <span class="pl-s1">p</span>.<span class="pl-s1">grad</span>)<span class="pl-c">#[0, 0, 0])</span>
            <span class="pl-s1">count</span> <span class="pl-c1">+=</span> <span class="pl-c1">1</span>
        <span class="pl-en">print</span>(<span class="pl-s">'num of params with grad: '</span>, <span class="pl-s1">count</span>)
    <span class="pl-k">else</span>:
        <span class="pl-s1">count</span> <span class="pl-c1">=</span> <span class="pl-c1">0</span>
        <span class="pl-k">for</span> <span class="pl-s1">n</span>, <span class="pl-s1">p</span> <span class="pl-c1">in</span> <span class="pl-s1">asr_model</span>.<span class="pl-s1">module</span>.<span class="pl-en">named_parameters</span>():
            <span class="pl-en">print</span>(<span class="pl-s">'grad of param {}, {}: '</span>.<span class="pl-en">format</span>(<span class="pl-s1">count</span>, <span class="pl-s1">n</span>), <span class="pl-s1">p</span>.<span class="pl-s1">grad</span>)<span class="pl-c">#[0, 0, 0])</span>
            <span class="pl-s1">count</span> <span class="pl-c1">+=</span> <span class="pl-c1">1</span>
        <span class="pl-en">print</span>(<span class="pl-s">'num of params with grad: '</span>, <span class="pl-s1">count</span>)


<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">'__main__'</span>:
    <span class="pl-s1">asr_model</span> <span class="pl-c1">=</span> <span class="pl-v">MultiHeadBLSTM</span>().<span class="pl-en">cuda</span>()
    <span class="pl-k">if</span> <span class="pl-en">len</span>(<span class="pl-s1">gpus</span>) <span class="pl-c1">&gt;</span> <span class="pl-c1">1</span>:
        <span class="pl-s1">asr_model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">DataParallel</span>(<span class="pl-s1">asr_model</span>, <span class="pl-s1">device_ids</span><span class="pl-c1">=</span><span class="pl-s1">gpus</span>)

    <span class="pl-s1">ctc_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">CTCLoss</span>(<span class="pl-s1">blank</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">reduction</span><span class="pl-c1">=</span><span class="pl-s">'mean'</span>, <span class="pl-s1">zero_infinity</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)

    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">rand</span>(<span class="pl-c1">4</span>, <span class="pl-c1">20</span>, <span class="pl-c1">20</span>)<span class="pl-c">#.cuda() </span>
    <span class="pl-s1">ilens</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">20</span>, <span class="pl-c1">20</span>, <span class="pl-c1">20</span>, <span class="pl-c1">20</span>])<span class="pl-c">#.cuda()</span>
    <span class="pl-s1">ys</span> <span class="pl-c1">=</span> [<span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>]),
          <span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>]),
          <span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>]),
          <span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>])]
    <span class="pl-s1">olens</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">3</span>, <span class="pl-c1">4</span>, <span class="pl-c1">5</span>, <span class="pl-c1">6</span>])<span class="pl-c">#.cuda()</span>

    <span class="pl-s1">y_true</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">cat</span>(<span class="pl-s1">ys</span>)<span class="pl-c">#.cuda()</span>

    <span class="pl-en">run_task</span>(<span class="pl-s1">asr_model</span>, <span class="pl-s1">ctc_loss</span>, <span class="pl-s1">x</span>, <span class="pl-s1">ilens</span>, <span class="pl-s1">y_true</span>, <span class="pl-s1">olens</span>)</pre></div>
<p dir="auto">If I run <code class="notranslate">python3 test.py 0</code> (gpu devices = [0]), everything is fine;<br>
however, if I run <code class="notranslate">python3 test.py 0,1</code> (gpu devices = [0, 1]), rnn receives no gradient.</p>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">All parameters of both single- and multi-gpu models should have non-zero gradients.</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.1
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Arch Linux
GCC version: (Arch Linux 9.3.0-1) 9.3.0
CMake version: version 3.17.0

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: GeForce GTX TITAN X
GPU 1: GeForce GTX 1060 6GB

Nvidia driver version: 440.64
cuDNN version: /usr/lib/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.2
[pip3] torch==1.4.1
[pip3] torchexp==0.1.0
[pip3] torchvision==0.5.0a0+c81ac87
[conda] Could not collect"><pre class="notranslate"><code class="notranslate">PyTorch version: 1.4.1
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Arch Linux
GCC version: (Arch Linux 9.3.0-1) 9.3.0
CMake version: version 3.17.0

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: GeForce GTX TITAN X
GPU 1: GeForce GTX 1060 6GB

Nvidia driver version: 440.64
cuDNN version: /usr/lib/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.2
[pip3] torch==1.4.1
[pip3] torchexp==0.1.0
[pip3] torchvision==0.5.0a0+c81ac87
[conda] Could not collect
</code></pre></div>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">I have tried running RNNs (LSTMs and GRUs) on multiple GPUs using <code class="notranslate">nn.DataParallel()</code> and when I print/plot the gradients, they are 0. They are not <code class="notranslate">Nan</code> or <code class="notranslate">None</code> but <code class="notranslate">0.0</code>. I have run the same code with CNNs and it worked fine. The RNN is part of <code class="notranslate">model.parameters()</code> so that is also not an issue.<br>
This issue exists only in 1.4.0. When I rolled back to 1.3.0 it worked fine.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<p dir="auto">Following is a self contained code. You can pass <code class="notranslate">multi=True</code> and <code class="notranslate">multi=False</code> to test the gradients on multiple or single GPU.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch

from torch import nn
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence

import warnings
warnings.filterwarnings(&quot;ignore&quot;)
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class BiLSTM(nn.Module):
    def __init__(self,):
        super(BiLSTM, self).__init__()
        self.lstm = nn.LSTM(10, 16, bidirectional = True, batch_first= True)
        self.classifier = nn.Linear(32, 1)

    def forward(self, embed, lengths):
        # sorted_len, sorted_idxs = torch.sort(lengths, descending =True)
        # embed = embed[ sorted_idxs ,: , :].to(device)
        sorted_len = lengths
        packed_embed = pack_padded_sequence(embed, sorted_len, batch_first = True, enforce_sorted=False).to(device)
        all_states, hidden_states = self.lstm(packed_embed)
        all_states, _ = pad_packed_sequence(all_states, batch_first = True)
        
        out = torch.cat((hidden_states[0][0], hidden_states[0][1]), dim = 1)
        out = self.classifier(out)
        return out.to(device)
    
def train(multi = False):
    seq_len = 5
    batch_size = 8
    input_dim = 10
    inp = torch.randn(batch_size, seq_len, input_dim).to(device)
    labels = torch.tensor([0,1,0,1,0,1,0,1])
    lens = torch.tensor([5,2,4,2,4,5,3,3]).to(device)
    
    model = BiLSTM().to(device)
    if multi:
        model = nn.DataParallel(model)
		
    optimizer = torch.optim.AdamW(model.parameters(), lr = 1e-4, weight_decay = 1e-3)
    criterion = nn.BCELoss()
    
    
    epochs = 1
    for i in range(epochs):
        preds = model(inp, lens)
        preds = F.sigmoid(preds)
        loss = criterion(preds,  labels.float().to(device))
        
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 10)
        optimizer.step()
        
        
        print(&quot;Gradients info..\n&quot;)
        model_log = model.module if multi else model
        for name, param in model_log.named_parameters():
            print(&quot;param {} grad = {}&quot;.format(name, param.grad.data.view(-1)))
        
    
if __name__ == '__main__':
    train(multi = False)"><pre class="notranslate"><code class="notranslate">import torch

from torch import nn
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence

import warnings
warnings.filterwarnings("ignore")
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class BiLSTM(nn.Module):
    def __init__(self,):
        super(BiLSTM, self).__init__()
        self.lstm = nn.LSTM(10, 16, bidirectional = True, batch_first= True)
        self.classifier = nn.Linear(32, 1)

    def forward(self, embed, lengths):
        # sorted_len, sorted_idxs = torch.sort(lengths, descending =True)
        # embed = embed[ sorted_idxs ,: , :].to(device)
        sorted_len = lengths
        packed_embed = pack_padded_sequence(embed, sorted_len, batch_first = True, enforce_sorted=False).to(device)
        all_states, hidden_states = self.lstm(packed_embed)
        all_states, _ = pad_packed_sequence(all_states, batch_first = True)
        
        out = torch.cat((hidden_states[0][0], hidden_states[0][1]), dim = 1)
        out = self.classifier(out)
        return out.to(device)
    
def train(multi = False):
    seq_len = 5
    batch_size = 8
    input_dim = 10
    inp = torch.randn(batch_size, seq_len, input_dim).to(device)
    labels = torch.tensor([0,1,0,1,0,1,0,1])
    lens = torch.tensor([5,2,4,2,4,5,3,3]).to(device)
    
    model = BiLSTM().to(device)
    if multi:
        model = nn.DataParallel(model)
		
    optimizer = torch.optim.AdamW(model.parameters(), lr = 1e-4, weight_decay = 1e-3)
    criterion = nn.BCELoss()
    
    
    epochs = 1
    for i in range(epochs):
        preds = model(inp, lens)
        preds = F.sigmoid(preds)
        loss = criterion(preds,  labels.float().to(device))
        
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 10)
        optimizer.step()
        
        
        print("Gradients info..\n")
        model_log = model.module if multi else model
        for name, param in model_log.named_parameters():
            print("param {} grad = {}".format(name, param.grad.data.view(-1)))
        
    
if __name__ == '__main__':
    train(multi = False)
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>

<h2 dir="auto">Environment</h2>
<p dir="auto">Please copy and paste the output from our<br>
<a href="https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py" rel="nofollow">environment collection script</a><br>
(or fill out the checklist below manually).</p>
<p dir="auto">You can get the script and run it with:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py"><pre class="notranslate"><code class="notranslate">wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
</code></pre></div>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): 1.4.0</li>
<li>OS (e.g., Linux): Linux</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): pip</li>
<li>Build command you used (if compiling from source): NA</li>
<li>Python version: 3.6.3</li>
<li>CUDA/cuDNN version: CUDA=10.0.13 / CuDNN = 7.4.2</li>
<li>GPU models and configuration: Nvidia 1080ti</li>
<li>Any other relevant information:</li>
</ul>
<h2 dir="auto">Additional context</h2>
<p dir="auto">NA</p>
<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></p>