xla::PjRtDevice* pjrt_dst_device =
      (*pjrt_dst_client)->LookupAddressableDevice(pjrt_dst_device_id).value();

  absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> buffer_or =
      src_device_buffer->CopyToDevice(pjrt_dst_device);
  if (!buffer_or.ok()) {
    done(buffer_or.status());
    return;
  }

  xla::PjRtBuffer* pjrt_buffer = (*buffer_or).get();

  if (static_cast<PjRtDeviceContext*>(recv_dev_context)

    if (absl::c_binary_search(constant_arg_indices, i)) {
      // Need to make sure it's on the host.
      inputs_storage.emplace_back(t->dtype(), t->shape());
      TF_RETURN_IF_ERROR(
          th->CopyToDevice(*context, /*d=*/nullptr, &inputs_storage.back()));
      inputs[i + offset] = &inputs_storage.back();
    } else {
      inputs[i + offset] = t;
    }
  }

  if (dev != nullptr) {