Lightning isn`t automatically moving a tensor to certain device

[StephennFernandes]

I am trying to implement this seq2seq cod from official pytorch docs here

I am using Lightning so i can scale the compute to TPU.

in a model there is an ouput tensor produced while decoding that collect a few samples of predictions on a single given input.
outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(device) where i set device as device = xm.xla_device()
here i have to explicitly do .to(device) because Lightning cannot detect it.

But in TPU as computations are distributed across the TPU cores. hence even after moving the tensor to one core doesn`t help
i need a way to sync the tensor to all the tpu cores while the computation is at it.

how to sync the tensor to all TPU cores

[awaelchli]

You can do

outputs = torch.zeros(max_len, batch_size, trg_vocab_size, device=self.device)

Can you show a code example of how you would like to "sync" this tensor?

[StephennFernandes]

You can checkout the link in mentioned above that's the exact code I used only with Pytorch Lighting.

On every single input data the seq2seq model creates this tensor with the given dimension that needs to be in the same Device as the device that's computing the prediction of the given input sample .

[awaelchli]

class Seq2Seq(LightningModule):
    def __init__(self,
                 encoder: nn.Module,
                 decoder: nn.Module):
        super().__init__()

        self.encoder = encoder
        self.decoder = decoder

        # delete this
        # self.device = device

    def forward(self,
                src: Tensor,
                trg: Tensor,
                teacher_forcing_ratio: float = 0.5) -> Tensor:

        batch_size = src.shape[1]
        max_len = trg.shape[0]
        trg_vocab_size = self.decoder.output_dim

        outputs = torch.zeros(max_len, batch_size, trg_vocab_size, device=self.device)

        encoder_outputs, hidden = self.encoder(src)

        # first input to the decoder is the <sos> token
        output = trg[0,:]

        for t in range(1, max_len):
            output, hidden = self.decoder(output, hidden, encoder_outputs)
            outputs[t] = output
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = output.max(1)[1]
            output = (trg[t] if teacher_force else top1)

        return outputs

here, just delete the device from the constructor and modify the line that creates the zeros tensor.