I am training a Sequence to Sequence LSTM model. The problem is that during the training loop, the loss calculation is complaining with the error: “Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!”.
So it seems like the computing the “score” from the model generates a tensor on the CPU instead of the GPU. I can fix this by just adjusting the line scores = model(data, targets) to  scores = model(data, targets).to(device), but that seems like an unnecessary passing of a tensor back and forth from the GPU to the CPU and then back to the GPU.
model = Seq2SeqPF(encoder_net, decoder_net).to(device)

load_from_checkpoint = False

if load_from_checkpoint:
    load_checkpoint(torch.load(os.path.join(CHECKPOINT_DIRECTORY, CHECKPOINT_NAME)),model, device)

for epoch in range(EPOCHS):
    print(f"Epoch: {epoch + 1}/{EPOCHS}")
    kbar = pkbar.Kbar(target=batches_per_epoch, width=8)

    if epoch % 5 == 0:
      checkpoint = {'state_dict': model.state_dict(), 
                    'optimizer': optimizer.state_dict()}
      save_checkpoint(checkpoint, 
                      CHECKPOINT_DIRECTORY, 
                      CHECKPOINT_NAME)

    for batch_idx, (data, targets) in enumerate(train_loader):
        
        data = data.to(device=device)
        targets = targets.to(device=device)

        # forward pass and compute error
        scores = model(data, targets) 
        loss = criterion(scores, targets)    # <---GENERATES ERROR ABOUT CPU AND GPU

    # backward pass and apply gradients.
        optimizer.zero_grad()
        loss.backward()

        # gradient descent step
        optimizer.step()

        kbar.update(batch_idx, values=[("loss", loss)])

The code seems standard. I also did push the model itself to the GPU device, and the encoder_net, and decoder_net are also network layers that are pushed to the GPU before the model code is pushed to the GPU.
Any suggestions on the right way to handle this?