Pytorch 如何为所有TPU设备更改一个TPU设备的变量？

我是否可以更改一个TPU设备内的变量，以便可以为其他TPU设备更改该变量

我尝试了以下方法：

!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.6-cp36-cp36m-linux_x86_64.whl
import torch
import torch_xla.core.xla_model as xm
a=torch.tensor([2], device=torch.device("cpu"))
def tr(rank, flags):
  device = xm.xla_device()  
  global a
  if xm.is_master_ordinal():
    a=torch.tensor([3], device=torch.device("cpu"))
  xm.rendezvous('download_only_once')
  b = a.to(device)
  print(b)

import torch_xla.distributed.xla_multiprocessing as xmp
flags={}
xmp.spawn(tr, args=(flags,), nprocs=8, start_method='fork')

它不起作用：其他设备的变量没有改变