File tree Expand file tree Collapse file tree 1 file changed +4
-2
lines changed
pylops_mpi/signalprocessing Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Original file line number Diff line number Diff line change @@ -128,7 +128,8 @@ def _matvec(self, x: DistributedArray) -> DistributedArray:
128128 for isl in range (self .nsls [self .rank ]):
129129 y1 [isl ] = ncp .dot (self .G [isl ], x [isl ])
130130 # gather results
131- y [:] = ncp .vstack (y ._allgather (y1 )).ravel ()
131+ y [:] = ncp .vstack (y ._allgather (y .base_comm , y .base_comm_nccl , y1 ,
132+ engine = y .engine )).ravel ()
132133 return y
133134
134135 def _rmatvec (self , x : NDArray ) -> NDArray :
@@ -165,5 +166,6 @@ def _rmatvec(self, x: NDArray) -> NDArray:
165166 y1 [isl ] = ncp .dot (x [isl ].T .conj (), self .G [isl ]).T .conj ()
166167
167168 # gather results
168- y [:] = ncp .vstack (y ._allgather (y1 )).ravel ()
169+ y [:] = ncp .vstack (y ._allgather (y .base_comm , y .base_comm_nccl , y1 ,
170+ engine = y .engine )).ravel ()
169171 return y
You can’t perform that action at this time.
0 commit comments