try to fix mpich compile error

CaRoLZhangxy · Apr 19, 2024 · 048c2af · 048c2af
1 parent 60605a9
commit 048c2af
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/source/op/pt/comm.cc b/source/op/pt/comm.cc
@@ -239,7 +239,13 @@ class Border : public torch::autograd::Function<Border> {
 #ifdef USE_MPI
   static void unpack_communicator(const torch::Tensor& communicator_tensor,
                                   MPI_Comm& mpi_comm) {
-    long int* communicator = communicator_tensor.data_ptr<long int>();
+
+#ifdef OMPI_MPI_H
+  long int* communicator = communicator_tensor.data_ptr<long int>();
+#else
+  long int* ptr = communicator_tensor.data_ptr<long int>();
+  int* communicator = reinterpret_cast<int*> (ptr);
+#endif
     mpi_comm = reinterpret_cast<MPI_Comm>(*communicator);
   }
 #endif