volcano.sh/volcano@v1.9.0/example/MindSpore-example/mindspore_gpu/gpu-test.py (about)

     1  import numpy as np
     2  import mindspore.context as context
     3  from mindspore import Tensor
     4  from mindspore.ops import functional as F
     5  from mindspore.communication.management import init, get_rank, get_group_size
     6  
     7  init('nccl')
     8  context.set_context(device_target="GPU")
     9  context.set_auto_parallel_context(parallel_mode="data_parallel", mirror_mean=True, device_num=get_group_size())
    10  
    11  x = Tensor(np.ones([1,3,3,4]).astype(np.float32))
    12  y = Tensor(np.ones([1,3,3,4]).astype(np.float32))
    13  print(F.tensor_add(x, y))