在本地计算机中无法正常运行使用 GPU 的 trax

Steps to reproduce:

1. Install trax

• pip install trax
• pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

1. Use jax Detect GPU

• code: import jax print(jax.devices())
• output: [gpu(id=0)]

Error logs:

1. Run the sample code of pre-trained Transformer in your Realme tutorial

• code: import os import numpy as np

  import trax

  Create a Transformer model.

  Pre-trained model config in gs://trax-ml/models/translation/ende_wmt32k.gin

  model = trax.models.Transformer( input_vocab_size=33300, d_model=512, d_ff=2048, n_heads=8, n_encoder_layers=6, n_decoder_layers=6, max_len=64, mode='predict')

  Initialize using pre-trained weights.

  model.init_from_file('gs://trax-ml/models/translation/ende_wmt32k.pkl.gz', weights_only=True) # input_signature=input_signature)

  Tokenize a sentence.

  sentence = 'It is nice to learn new things today!' tokenized = list(trax.data.tokenize(iter([sentence]), # Operates on streams. vocab_dir='gs://trax-ml/vocabs/', vocab_file='ende_32k.subword'))[0]

  Decode from the Transformer.

  tokenized = tokenized[None, :] # Add batch dimension. tokenized_translation = trax.supervised.decoding.autoregressive_sample( model, tokenized, temperature=0.0) # Higher temperature: more diverse results.

  De-tokenize,

  tokenized_translation = tokenized_translation[0][:-1] # Remove batch and EOS. translation = trax.data.detokenize(tokenized_translation, vocab_dir='gs://trax-ml/vocabs/', vocab_file='ende_32k.subword') print(translation)

• Error Output: 2023-06-22 15:58:35.266959: W tensorflow/tsl/platform/cloud/google_auth_provider.cc:184] All attempts to get a Google authentication bearer token failed, returning an