diff --git a/src/KOKKOS/kokkos.cpp b/src/KOKKOS/kokkos.cpp index cc53ff72e6..10e64f0bd6 100644 --- a/src/KOKKOS/kokkos.cpp +++ b/src/KOKKOS/kokkos.cpp @@ -160,6 +160,12 @@ KokkosLMP::KokkosLMP(LAMMPS *lmp, int narg, char **arg) : Pointers(lmp) if (device >= skip_gpu) device++; set_flag = 1; } + if ((str = getenv("PMI_LOCAL_RANK"))) { + int local_rank = atoi(str); + device = local_rank % ngpus; + if (device >= skip_gpu) device++; + set_flag = 1; + } if (ngpus > 1 && !set_flag) error->all(FLERR,"Could not determine local MPI rank for multiple "