I am using pytorch xla with TPU on google colab. As expected computations are indeed faster, however when I go back to cpu to save my tensors, the “action” of going back to the cpu takes a lot of time for big tensors. Here is the function I run in parallel:
def _mp_fn(index, batch_size, model, typee, directory_join):
    device = xm.xla_device()
    xm.master_print(f"Here")
    if not xm.is_master_ordinal():
        xm.rendezvous("download_only_once")

    loader = sherBERTEmbedding(f_df, typee)

    if xm.is_master_ordinal():
        xm.rendezvous("download_only_once")
    xm.master_print(f"There")

    sampler = DistributedSampler(loader, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal(), shuffle=False)
    dataloader = DataLoader(
        loader,
        batch_size=batch_size,
        sampler=sampler,
        # drop_last=True,
    )

    model = model.to(device)
    model.eval()

    length = len(dataloader)

    para_loader = pl.ParallelLoader(dataloader, [device])
    num_b = len(dataloader)
    tensor_embeddings = torch.zeros((batch_size*num_b, 153, 768), dtype=torch.float).to(device)
    tensor_ids = torch.zeros((batch_size*num_b, 153), dtype=torch.long).to(device)
    tensor_index = torch.zeros(batch_size*num_b, dtype=torch.long).to(device)

    toremove = 0
    prev_ind = 0

    with torch.no_grad():
        for t, (x, x_pad_mask, index) in tqdm(
            enumerate(para_loader.per_device_loader(device)), total=length
        ):
            x = x.to(device=device, dtype=torch.long)
            x_pad_mask = x_pad_mask.to(device=device, dtype=torch.long)
            result = get_embeddings_best_token(x, x_pad_mask, model_embed, index_best_tokens)
            ids = get_ids_best_tokens(x, index_best_tokens)

            tensor_embeddings[prev_ind: prev_ind + result.size(0), :, :] = result
            tensor_ids[prev_ind: prev_ind + ids.size(0), :] = ids
            tensor_index[prev_ind: prev_ind +index.size(0)] = index.squeeze(1)
            prev_ind += result.size(0)


            if result.size(0) != batch_size:
                toremove = (batch_size - result.size(0))

        if toremove == 0:
            toremove = tensor_embeddings.size(0)+1


        #saving
        directory_embeddings = os.path.join(os.path.join(directory_join, typee), "embeddings")
        directory_ids = os.path.join(os.path.join(directory_join, typee), "ids")
        directory_index = os.path.join(os.path.join(directory_join, typee), "index")

        tosave = tensor_index[:-toremove]
        ltosave = tosave.cpu()
        torch.save(tosave, os.path.join(directory_index, str(xm.get_ordinal())+".p"))
        print("Saving 1")

        tosave = tensor_ids[:-toremove, :].squeeze()
        ltosave = tosave.cpu()
        torch.save(tosave, os.path.join(directory_ids, str(xm.get_ordinal())+".p"))
        print("Saving 2")

        tosave = tensor_embeddings[:-toremove, :, : ].view(-1,768).detach()
        ltosave = tosave.cpu()
        print("Saving 3")
        torch.save(ltosave, os.path.join(directory_embeddings, str(xm.get_ordinal())+".p"))





def get_embeddings_best_token(root, mask, model, index_best_tokens):
    res = model((root, mask))
    return res[:,index_best_tokens[0]:index_best_tokens[-1]+1, :]

def get_ids_best_tokens(x, index_best_tokens):
    return x[:,index_best_tokens[0]:index_best_tokens[-1]+1]

And I run in parallel with this:
t1 = time.time()
get_embeddings(8, model_embed, "myType", JOIN_DIR)
t2 = time.time()

I know that going back to cpu between computations on TPU is a bottleneck, however here I am just trying to save the tensors after the “computations”. I know that a special function exists to save tensors:  torch_xla.core.xla_model.save(data, file_or_path, master_only=True, global_master=False) . However, they explain that basically, it converts tensors from xla tensors to cpu tensors and then saves them. When I use this function, it takes around the same amount of time.