Repository: espnet/espnet
Issue: Multi-GPU training failure with PyTorch 1.5.1

E.g. in `egs/libritts/tts1`, `$ CUDA_VISIBLE_DEVICES=0,1 ./run.sh --ngpu 2 --trans_type phn` will result in the following error in the log file,
```
Exception in main training loop: Caught StopIteration in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 60, in _worker
    output = module(*input, **kwargs)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/e2e_tts_tacotron2.py", line 746, in forward
    after_outs, before_outs, logits, att_ws = self.dec(hs, hlens, ys)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/tacotron2/decoder.py", line 436, in forward
    att_c, att_w = self.att(hs, hlens, z_list[0], prev_att_w)                                                                                                                                           File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)                                                                                                                                                             File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/rnn/attentions.py", line 365, in forward
    self.mask = to_device(self, make_pad_mask(enc_hs_len))
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/nets_utils.py", line 24, in to_device
    device = next(m.parameters()).device
StopIteration
Traceback (most recent call last):                                                                                                                                                                      File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/chainer/training/trainer.py", line 316, in run
    update()
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/tts/pytorch_backend/tts.py", line 171, in update                                                                                           self.update_core()
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/tts/pytorch_backend/tts.py", line 151, in update_core
    loss = self.model(**x).mean() / self.accum_grad                                                                                                                                                     File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 155, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 165, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in parallel_apply
    output.reraise()
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/_utils.py", line 395, in reraise
    raise self.exc_type(msg)
Will finalize trainer extensions and updater before reraising the exception.
Traceback (most recent call last):
  File "/home/huang18/VC/Experiments/espnet/egs/ljspeech/tts1/../../../espnet/bin/tts_train.py", line 355, in <module>
    main(sys.argv[1:])
  File "/home/huang18/VC/Experiments/espnet/egs/ljspeech/tts1/../../../espnet/bin/tts_train.py", line 349, in main
    train(args)
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/tts/pytorch_backend/tts.py", line 563, in train
    trainer.run()
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/chainer/training/trainer.py", line 349, in run
    six.reraise(*exc_info)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/six.py", line 693, in reraise
    raise value
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/chainer/training/trainer.py", line 316, in run
    update()
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/tts/pytorch_backend/tts.py", line 171, in update
    self.update_core()
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/tts/pytorch_backend/tts.py", line 151, in update_core
    loss = self.model(**x).mean() / self.accum_grad
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 155, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 165, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in parallel_apply
    output.reraise()
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/_utils.py", line 395, in reraise
    raise self.exc_type(msg)
StopIteration: Caught StopIteration in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 60, in _worker
    output = module(*input, **kwargs)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/e2e_tts_tacotron2.py", line 746, in forward
    after_outs, before_outs, logits, att_ws = self.dec(hs, hlens, ys)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/tacotron2/decoder.py", line 436, in forward
    att_c, att_w = self.att(hs, hlens, z_list[0], prev_att_w)
  File "/home/huang18/VC/Experiments/espnet/tools/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/rnn/attentions.py", line 365, in forward
    self.mask = to_device(self, make_pad_mask(enc_hs_len))
  File "/nas01/internal/wenchin-h/VC/Experiments/espnet/espnet/nets/pytorch_backend/nets_utils.py", line 24, in to_device
    device = next(m.parameters()).device
StopIteration
```

According to this issue: https://github.com/pytorch/pytorch/issues/40457 , a quick fix is to use PyTorch 1.4.0.
But as the answer states, https://github.com/pytorch/pytorch/issues/40457#issuecomment-648396469 , accessing model parameters with `parameter()`  is the root cause. 
For example, in
https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/nets_utils.py#L24 .
I am not sure if there exists a small, fast solution to this issue that does not touch this seemingly fundamental, highly used function... Just reporting it in case someone else is also having this problem.

Changed files:
espnet/lm/pytorch_backend/extlm.py, espnet/nets/pytorch_backend/ctc.py, espnet/nets/pytorch_backend/rnn/attentions.py, espnet/nets/pytorch_backend/rnn/decoders.py, espnet/nets/pytorch_backend/nets_utils.py, espnet/nets/pytorch_backend/lm/default.py, espnet/nets/pytorch_backend/rnn/encoders.py, espnet/nets/pytorch_backend/e2e_asr_mix.py

diff --git a/espnet/lm/pytorch_backend/extlm.py b/espnet/lm/pytorch_backend/extlm.py
index da8999cb9..557487092 100644
--- a/espnet/lm/pytorch_backend/extlm.py
+++ b/espnet/lm/pytorch_backend/extlm.py
@@ -50,8 +50,8 @@ class MultiLevelLM(nn.Module):
     def forward(self, state, x):
         # update state with input label x
         if state is None:  # make initial states and log-prob vectors
-            self.var_word_eos = to_device(self, self.var_word_eos)
-            self.var_word_unk = to_device(self, self.var_word_eos)
+            self.var_word_eos = to_device(x, self.var_word_eos)
+            self.var_word_unk = to_device(x, self.var_word_eos)
             wlm_state, z_wlm = self.wordlm(None, self.var_word_eos)
             wlm_logprobs = F.log_softmax(z_wlm, dim=1)
             clm_state, z_clm = self.subwordlm(None, x)
@@ -64,7 +64,7 @@ class MultiLevelLM(nn.Module):
             xi = int(x)
             if xi == self.space:  # inter-word transition
                 if node is not None and node[1] >= 0:  # check if the node is word end
-                    w = to_device(self, torch.LongTensor([node[1]]))
+                    w = to_device(x, torch.LongTensor([node[1]]))
                 else:  # this node is not a word end, which means <unk>
                     w = self.var_word_unk
                 # update wordlm state and log-prob vector
@@ -80,7 +80,7 @@ class MultiLevelLM(nn.Module):
                 clm_logprob += log_y[0, xi]
             else:  # if open_vocab flag is disabled, return 0 probabilities
                 log_y = to_device(
-                    self, torch.full((1, self.subword_dict_size), self.logzero)
+                    x, torch.full((1, self.subword_dict_size), self.logzero)
                 )
                 return (clm_state, wlm_state, wlm_logprobs, None, log_y, 0.0), log_y
 
@@ -107,7 +107,7 @@ class MultiLevelLM(nn.Module):
     def final(self, state):
         clm_state, wlm_state, wlm_logprobs, node, log_y, clm_logprob = state
         if node is not None and node[1] >= 0:  # check if the node is word end
-            w = to_device(self, torch.LongTensor([node[1]]))
+            w = to_device(wlm_logprobs, torch.LongTensor([node[1]]))
         else:  # this node is not a word end, which means <unk>
             w = self.var_word_unk
         wlm_state, z_wlm = self.wordlm(wlm_state, w)
@@ -140,9 +140,9 @@ class LookAheadWordLM(nn.Module):
     def forward(self, state, x):
         # update state with input label x
         if state is None:  # make initial states and cumlative probability vector
-            self.var_word_eos = to_device(self, self.var_word_eos)
-            self.var_word_unk = to_device(self, self.var_word_eos)
-            self.zero_tensor = to_device(self, self.zero_tensor)
+            self.var_word_eos = to_device(x, self.var_word_eos)
+            self.var_word_unk = to_device(x, self.var_word_eos)
+            self.zero_tensor = to_device(x, self.zero_tensor)
             wlm_state, z_wlm = self.wordlm(None, self.var_word_eos)
             cumsum_probs = torch.cumsum(F.softmax(z_wlm, dim=1), dim=1)
             new_node = self.lexroot
@@ -152,7 +152,7 @@ class LookAheadWordLM(nn.Module):
             xi = int(x)
             if xi == self.space:  # inter-word transition
                 if node is not None and node[1] >= 0:  # check if the node is word end
-                    w = to_device(self, torch.LongTensor([node[1]]))
+                    w = to_device(x, torch.LongTensor([node[1]]))
                 else:  # this node is not a word end, which means <unk>
                     w = self.var_word_unk
                 # update wordlm state and cumlative probability vector
@@ -165,7 +165,7 @@ class LookAheadWordLM(nn.Module):
                 new_node = None
             else:  # if open_vocab flag is disabled, return 0 probabilities
                 log_y = to_device(
-                    self, torch.full((1, self.subword_dict_size), self.logzero)
+                    x, torch.full((1, self.subword_dict_size), self.logzero)
                 )
                 return (wlm_state, None, None), log_y
 
@@ -179,7 +179,7 @@ class LookAheadWordLM(nn.Module):
             )
             if sum_prob < self.zero:
                 log_y = to_device(
-                    self, torch.full((1, self.subword_dict_size), self.logzero)
+                    x, torch.full((1, self.subword_dict_size), self.logzero)
                 )
                 return (wlm_state, cumsum_probs, new_node), log_y
             # set <unk> probability as a default value
@@ -187,7 +187,7 @@ class LookAheadWordLM(nn.Module):
                 cumsum_probs[:, self.word_unk] - cumsum_probs[:, self.word_unk - 1]
             )
             y = to_device(
-                self,
+                x,
                 torch.full(
                     (1, self.subword_dict_size), float(unk_prob) * self.oov_penalty
                 ),
@@ -207,13 +207,13 @@ class LookAheadWordLM(nn.Module):
                 y[:, self.eos] = self.zero
             log_y = torch.log(torch.max(y, self.zero_tensor))  # clip to avoid log(0)
         else:  # if no path in the tree, transition probability is one
-            log_y = to_device(self, torch.zeros(1, self.subword_dict_size))
+            log_y = to_device(x, torch.zeros(1, self.subword_dict_size))
         return (wlm_state, cumsum_probs, new_node), log_y
 
     def final(self, state):
         wlm_state, cumsum_probs, node = state
         if node is not None and node[1] >= 0:  # check if the node is word end
-            w = to_device(self, torch.LongTensor([node[1]]))
+            w = to_device(cumsum_probs, torch.LongTensor([node[1]]))
         else:  # this node is not a word end, which means <unk>
             w = self.var_word_unk
         wlm_state, z_wlm = self.wordlm(wlm_state, w)
diff --git a/espnet/nets/pytorch_backend/ctc.py b/espnet/nets/pytorch_backend/ctc.py
index 7fe426940..dc27caece 100644
--- a/espnet/nets/pytorch_backend/ctc.py
+++ b/espnet/nets/pytorch_backend/ctc.py
@@ -109,8 +109,8 @@ class CTC(torch.nn.Module):
             ys_hat = ys_hat.to(dtype=torch.float32)
         if self.ctc_type == "builtin":
             # use GPU when using the cuDNN implementation
-            ys_true = to_device(self, ys_true)
-        self.loss = to_device(self, self.loss_fn(ys_hat, ys_true, hlens, olens)).to(
+            ys_true = to_device(hs_pad, ys_true)
+        self.loss = to_device(hs_pad, self.loss_fn(ys_hat, ys_true, hlens, olens)).to(
             dtype=dtype
         )
         if self.reduce:
diff --git a/espnet/nets/pytorch_backend/e2e_asr_mix.py b/espnet/nets/pytorch_backend/e2e_asr_mix.py
index b874a5cfa..fccd84c8d 100644
--- a/espnet/nets/pytorch_backend/e2e_asr_mix.py
+++ b/espnet/nets/pytorch_backend/e2e_asr_mix.py
@@ -799,7 +799,7 @@ class EncoderMix(torch.nn.Module):
                 xs_pad_sd[ns], ilens_sd[ns], _ = module(xs_pad_sd[ns], ilens_sd[ns])
 
         # make mask to remove bias value in padded part
-        mask = to_device(self, make_pad_mask(ilens_sd[0]).unsqueeze(-1))
+        mask = to_device(xs_pad, make_pad_mask(ilens_sd[0]).unsqueeze(-1))
 
         return [x.masked_fill(mask, 0.0) for x in xs_pad_sd], ilens_sd, None
 
diff --git a/espnet/nets/pytorch_backend/lm/default.py b/espnet/nets/pytorch_backend/lm/default.py
index 506df1e8e..1c41af2c6 100644
--- a/espnet/nets/pytorch_backend/lm/default.py
+++ b/espnet/nets/pytorch_backend/lm/default.py
@@ -344,14 +344,11 @@ class RNNLM(nn.Module):
     def forward(self, state, x):
         """Forward neural networks."""
         if state is None:
-            h = [
-                to_device(self, self.zero_state(x.size(0)))
-                for n in range(self.n_layers)
-            ]
+            h = [to_device(x, self.zero_state(x.size(0))) for n in range(self.n_layers)]
             state = {"h": h}
             if self.typ == "lstm":
                 c = [
-                    to_device(self, self.zero_state(x.size(0)))
+                    to_device(x, self.zero_state(x.size(0)))
                     for n in range(self.n_layers)
                 ]
                 state = {"c": c, "h": h}
diff --git a/espnet/nets/pytorch_backend/nets_utils.py b/espnet/nets/pytorch_backend/nets_utils.py
index 3f5b41ffb..b0e64188b 100644
--- a/espnet/nets/pytorch_backend/nets_utils.py
+++ b/espnet/nets/pytorch_backend/nets_utils.py
@@ -20,8 +20,14 @@ def to_device(m, x):
         Tensor: Torch tensor located in the same place as torch module.
 
     """
-    assert isinstance(m, torch.nn.Module)
-    device = next(m.parameters()).device
+    if isinstance(m, torch.nn.Module):
+        device = next(m.parameters()).device
+    elif isinstance(m, torch.Tensor):
+        device = m.device
+    else:
+        raise TypeError(
+            "Expected torch.nn.Module or torch.tensor, " f"bot got: {type(m)}"
+        )
     return x.to(device)
 
 
diff --git a/espnet/nets/pytorch_backend/rnn/attentions.py b/espnet/nets/pytorch_backend/rnn/attentions.py
index 8c37af2e3..8458c156f 100644
--- a/espnet/nets/pytorch_backend/rnn/attentions.py
+++ b/espnet/nets/pytorch_backend/rnn/attentions.py
@@ -157,7 +157,7 @@ class AttDot(torch.nn.Module):
 
         # NOTE consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
         w = F.softmax(scaling * e, dim=1)
 
@@ -235,7 +235,7 @@ class AttAdd(torch.nn.Module):
 
         # NOTE consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
         w = F.softmax(scaling * e, dim=1)
 
@@ -362,7 +362,7 @@ class AttLoc(torch.nn.Module):
 
         # NOTE: consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
 
         # apply monotonic attention constraint (mainly for TTS)
@@ -446,7 +446,9 @@ class AttCov(torch.nn.Module):
         # initialize attention weight with uniform dist.
         if att_prev_list is None:
             # if no bias, 0 0-pad goes 0
-            att_prev_list = to_device(self, (1.0 - make_pad_mask(enc_hs_len).float()))
+            att_prev_list = to_device(
+                enc_hs_pad, (1.0 - make_pad_mask(enc_hs_len).float())
+            )
             att_prev_list = [
                 att_prev_list / att_prev_list.new(enc_hs_len).unsqueeze(-1)
             ]
@@ -467,7 +469,7 @@ class AttCov(torch.nn.Module):
 
         # NOTE consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
         w = F.softmax(scaling * e, dim=1)
         att_prev_list += [w]
@@ -562,7 +564,7 @@ class AttLoc2D(torch.nn.Module):
         if att_prev is None:
             # B * [Li x att_win]
             # if no bias, 0 0-pad goes 0
-            att_prev = to_device(self, (1.0 - make_pad_mask(enc_hs_len).float()))
+            att_prev = to_device(enc_hs_pad, (1.0 - make_pad_mask(enc_hs_len).float()))
             att_prev = att_prev / att_prev.new(enc_hs_len).unsqueeze(-1)
             att_prev = att_prev.unsqueeze(1).expand(-1, self.att_win, -1)
 
@@ -584,7 +586,7 @@ class AttLoc2D(torch.nn.Module):
 
         # NOTE consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
         w = F.softmax(scaling * e, dim=1)
 
@@ -681,7 +683,7 @@ class AttLocRec(torch.nn.Module):
         if att_prev_states is None:
             # initialize attention weight with uniform dist.
             # if no bias, 0 0-pad goes 0
-            att_prev = to_device(self, (1.0 - make_pad_mask(enc_hs_len).float()))
+            att_prev = to_device(enc_hs_pad, (1.0 - make_pad_mask(enc_hs_len).float()))
             att_prev = att_prev / att_prev.new(enc_hs_len).unsqueeze(-1)
 
             # initialize lstm states
@@ -712,7 +714,7 @@ class AttLocRec(torch.nn.Module):
 
         # NOTE consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
         w = F.softmax(scaling * e, dim=1)
 
@@ -802,7 +804,9 @@ class AttCovLoc(torch.nn.Module):
         if att_prev_list is None:
             # if no bias, 0 0-pad goes 0
             mask = 1.0 - make_pad_mask(enc_hs_len).float()
-            att_prev_list = [to_device(self, mask / mask.new(enc_hs_len).unsqueeze(-1))]
+            att_prev_list = [
+                to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1))
+            ]
 
         # att_prev_list: L' * [B x T] => cov_vec B x T
         cov_vec = sum(att_prev_list)
@@ -825,7 +829,7 @@ class AttCovLoc(torch.nn.Module):
 
         # NOTE consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
         w = F.softmax(scaling * e, dim=1)
         att_prev_list += [w]
@@ -932,7 +936,7 @@ class AttMultiHeadDot(torch.nn.Module):
 
             # NOTE consider zero padding when compute w.
             if self.mask is None:
-                self.mask = to_device(self, make_pad_mask(enc_hs_len))
+                self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
             e.masked_fill_(self.mask, -float("inf"))
             w += [F.softmax(self.scaling * e, dim=1)]
 
@@ -1049,7 +1053,7 @@ class AttMultiHeadAdd(torch.nn.Module):
 
             # NOTE consider zero padding when compute w.
             if self.mask is None:
-                self.mask = to_device(self, make_pad_mask(enc_hs_len))
+                self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
             e.masked_fill_(self.mask, -float("inf"))
             w += [F.softmax(self.scaling * e, dim=1)]
 
@@ -1185,7 +1189,9 @@ class AttMultiHeadLoc(torch.nn.Module):
             for _ in six.moves.range(self.aheads):
                 # if no bias, 0 0-pad goes 0
                 mask = 1.0 - make_pad_mask(enc_hs_len).float()
-                att_prev += [to_device(self, mask / mask.new(enc_hs_len).unsqueeze(-1))]
+                att_prev += [
+                    to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1))
+                ]
 
         c = []
         w = []
@@ -1204,7 +1210,7 @@ class AttMultiHeadLoc(torch.nn.Module):
 
             # NOTE consider zero padding when compute w.
             if self.mask is None:
-                self.mask = to_device(self, make_pad_mask(enc_hs_len))
+                self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
             e.masked_fill_(self.mask, -float("inf"))
             w += [F.softmax(scaling * e, dim=1)]
 
@@ -1339,7 +1345,9 @@ class AttMultiHeadMultiResLoc(torch.nn.Module):
             for _ in six.moves.range(self.aheads):
                 # if no bias, 0 0-pad goes 0
                 mask = 1.0 - make_pad_mask(enc_hs_len).float()
-                att_prev += [to_device(self, mask / mask.new(enc_hs_len).unsqueeze(-1))]
+                att_prev += [
+                    to_device(enc_hs_pad, mask / mask.new(enc_hs_len).unsqueeze(-1))
+                ]
 
         c = []
         w = []
@@ -1358,7 +1366,7 @@ class AttMultiHeadMultiResLoc(torch.nn.Module):
 
             # NOTE consider zero padding when compute w.
             if self.mask is None:
-                self.mask = to_device(self, make_pad_mask(enc_hs_len))
+                self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
             e.masked_fill_(self.mask, -float("inf"))
             w += [F.softmax(self.scaling * e, dim=1)]
 
@@ -1482,7 +1490,7 @@ class AttForward(torch.nn.Module):
 
         # NOTE: consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
 
         # apply monotonic attention constraint (mainly for TTS)
@@ -1617,7 +1625,7 @@ class AttForwardTA(torch.nn.Module):
 
         # NOTE consider zero padding when compute w.
         if self.mask is None:
-            self.mask = to_device(self, make_pad_mask(enc_hs_len))
+            self.mask = to_device(enc_hs_pad, make_pad_mask(enc_hs_len))
         e.masked_fill_(self.mask, -float("inf"))
 
         # apply monotonic attention constraint (mainly for TTS)
diff --git a/espnet/nets/pytorch_backend/rnn/decoders.py b/espnet/nets/pytorch_backend/rnn/decoders.py
index 2250d3bbd..8f1b22028 100644
--- a/espnet/nets/pytorch_backend/rnn/decoders.py
+++ b/espnet/nets/pytorch_backend/rnn/decoders.py
@@ -250,7 +250,7 @@ class Decoder(torch.nn.Module, ScorerInterface):
                 logging.info(" scheduled sampling ")
                 z_out = self.output(z_all[-1])
                 z_out = np.argmax(z_out.detach().cpu(), axis=1)
-                z_out = self.dropout_emb(self.embed(to_device(self, z_out)))
+                z_out = self.dropout_emb(self.embed(to_device(hs_pad[0], z_out)))
                 ey = torch.cat((z_out, att_c), dim=1)  # utt x (zdim + hdim)
             else:
                 ey = torch.cat((eys[:, i, :], att_c), dim=1)  # utt x (zdim + hdim)
@@ -302,7 +302,7 @@ class Decoder(torch.nn.Module, ScorerInterface):
 
         if self.labeldist is not None:
             if self.vlabeldist is None:
-                self.vlabeldist = to_device(self, torch.from_numpy(self.labeldist))
+                self.vlabeldist = to_device(hs_pad[0], torch.from_numpy(self.labeldist))
             loss_reg = -torch.sum(
                 (F.log_softmax(y_all, dim=1) * self.vlabeldist).view(-1), dim=0
             ) / len(ys_in)
@@ -681,7 +681,7 @@ class Decoder(torch.nn.Module, ScorerInterface):
             weights_ctc_dec = [1.0]
 
         n_bb = batch * beam
-        pad_b = to_device(self, torch.arange(batch) * beam).view(-1, 1)
+        pad_b = to_device(h[0], torch.arange(batch) * beam).view(-1, 1)
 
         max_hlen = np.amin([max(hlens[idx]) for idx in range(self.num_encs)])
         if recog_args.maxlenratio == 0:
@@ -694,18 +694,18 @@ class Decoder(torch.nn.Module, ScorerInterface):
 
         # initialization
         c_prev = [
-            to_device(self, torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
+            to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
         ]
         z_prev = [
-            to_device(self, torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
+            to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
         ]
         c_list = [
-            to_device(self, torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
+            to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
         ]
         z_list = [
-            to_device(self, torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
+            to_device(h[0], torch.zeros(n_bb, self.dunits)) for _ in range(self.dlayers)
         ]
-        vscores = to_device(self, torch.zeros(batch, beam))
+        vscores = to_device(h[0], torch.zeros(batch, beam))
 
         rnnlm_state = None
         if self.num_encs == 1:
@@ -776,7 +776,7 @@ class Decoder(torch.nn.Module, ScorerInterface):
         for i in six.moves.range(maxlen):
             logging.debug("position " + str(i))
 
-            vy = to_device(self, torch.LongTensor(self._get_last_yseq(yseq)))
+            vy = to_device(h[0], torch.LongTensor(self._get_last_yseq(yseq)))
             ey = self.dropout_emb(self.embed(vy))
             if self.num_encs == 1:
                 att_c, att_w = self.att[att_idx](
@@ -857,7 +857,7 @@ class Decoder(torch.nn.Module, ScorerInterface):
             yseq = self._index_select_list(yseq, accum_padded_beam_ids)
             yseq = self._append_ids(yseq, accum_odim_ids)
             vscores = accum_best_scores
-            vidx = to_device(self, torch.LongTensor(accum_padded_beam_ids))
+            vidx = to_device(h[0], torch.LongTensor(accum_padded_beam_ids))
 
             a_prev = []
             num_atts = self.num_encs if self.num_encs == 1 else self.num_encs + 1
diff --git a/espnet/nets/pytorch_backend/rnn/encoders.py b/espnet/nets/pytorch_backend/rnn/encoders.py
index 370192f57..d388a5479 100644
--- a/espnet/nets/pytorch_backend/rnn/encoders.py
+++ b/espnet/nets/pytorch_backend/rnn/encoders.py
@@ -317,7 +317,7 @@ class Encoder(torch.nn.Module):
             current_states.append(states)
 
         # make mask to remove bias value in padded part
-        mask = to_device(self, make_pad_mask(ilens).unsqueeze(-1))
+        mask = to_device(xs_pad, make_pad_mask(ilens).unsqueeze(-1))
 
         return xs_pad.masked_fill(mask, 0.0), ilens, current_states