We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 6bfd2e1 commit 6276d4eCopy full SHA for 6276d4e
pylops_mpi/utils/_nccl.py
@@ -110,11 +110,11 @@ def initialize_nccl_comm() -> nccl.NcclCommunicator:
110
111
# Create a communicator for ranks on the same node
112
node_comm = comm.Split_type(MPI.COMM_TYPE_SHARED)
113
- ranks_on_node = node_comm.Get_size()
+ size_node = node_comm.Get_size()
114
115
device_id = int(
116
os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
117
- or (rank % ranks_on_node) % cp.cuda.runtime.getDeviceCount()
+ or (rank % size_node) % cp.cuda.runtime.getDeviceCount()
118
)
119
cp.cuda.Device(device_id).use()
120
0 commit comments