Dear Pytorch team,
Greetings!
I am trying to run a model with multiple GPUs, but an arguments are located on different GPUs error occurs. I researched this problem in pytorch forum but all the solutions seem not to work.
Particularly, I have been trying to use register_buffer
self.register_buffer("positions", positions)

to solve the problem
RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorMasked.cu:35

But the same error still occurs.
So my torch vision is 1.0.0, and I have referred to this post 6 to try to solve. But it didn’t work.
My code is below. It is implementing the positional embedding of transformer model. The error occurs at the very end of this piece.
class SinusoidalPositionalEmbedding(nn.Module):
    """This module produces sinusoidal positional embeddings of any length.
    Padding symbols are ignored, but it is necessary to specify whether padding
    is added on the left side (left_pad=True) or right side (left_pad=False).
    """

    def __init__(self, embedding_dim, padding_idx=0, left_pad=0, init_size=128):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx
        self.left_pad = left_pad
        self.weights = SinusoidalPositionalEmbedding.get_embedding(
            init_size,
            embedding_dim,
            padding_idx,
        )
        # Here is the buffers
        self.register_buffer('_float_tensor', torch.FloatTensor(1))
        positions = None
        self.register_buffer("positions", positions)
        mask = None
        self.register_buffer("mask", mask)
    @staticmethod
    def get_embedding(num_embeddings, embedding_dim, padding_idx=None):
        """Build sinusoidal embeddings.
        This matches the implementation in tensor2tensor, but differs slightly
        from the description in Section 3.5 of "Attention Is All You Need".
        """
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
        if embedding_dim % 2 == 1:
            # zero pad
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        if padding_idx is not None:
            emb[padding_idx, :] = 0
        return emb

    def forward(self, input):
        """Input is expected to be of size [bsz x seqlen]."""
        bsz, seq_len = input.size()
        max_pos = self.padding_idx + 1 + seq_len
        if self.weights is None or max_pos > self.weights.size(0):
            self.weights = SinusoidalPositionalEmbedding.get_embedding(
                max_pos,
                self.embedding_dim,
                self.padding_idx,
            )
        self.weights = self.weights.type_as(self._float_tensor)
        max_pos = self.padding_idx + 1 + input.size(1)
        if not hasattr(make_positions, 'range_buf'):
            make_positions.range_buf = input.new()
        make_positions.range_buf = make_positions.range_buf.type_as(input)
        if make_positions.range_buf.numel() < max_pos:
            torch.arange(self.padding_idx + 1, max_pos, out=make_positions.range_buf)
        self.mask = input.ne(self.padding_idx)
        self.positions = make_positions.range_buf[:input.size(1)].expand_as(input)
        if self.left_pad:
            positions = positions - self.mask.size(1) + self.mask.long().sum(dim=1).unsqueeze(1)

        # !!! Here is where the error points to
        self.positions = input.clone().masked_scatter_(self.mask, self.positions[self.mask]).long()
        return self.weights.index_select(0, self.positions.view(-1)).view(bsz, seq_len, -1).detach()


The error I get:
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/qianlim/dl_signal/transformer/modules/position_embedding.py", line 92, in forward
    self.positions = input.clone().masked_scatter_(self.mask, self.positions[self.mask]).long()
RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorMasked.cu:35


Thank you so much in advance!