fix

Tonny-Gu · Mar 11, 2022 · c58a495 · c58a495
1 parent bd2c3d7
commit c58a495
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 1 deletion.
diff --git a/include/raf/communicator.h b/include/raf/communicator.h
@@ -146,7 +146,7 @@ class CommunicatorPool {
     comm_.clear();
   }
 
- public:
+ private:
   std::map<CommunicatorID, Communicator> comm_;
 };
 

diff --git a/src/distributed/common/communicator.cc b/src/distributed/common/communicator.cc
@@ -121,6 +121,8 @@ RAF_REGISTER_GLOBAL("raf.distributed.RemoveCommunicator").set_body_typed([]() {
   CommunicatorPool::Get()->Remove();
 });
 
+RAF_REGISTER_OBJECT_REFLECT(VoidCommunicatorObj);
+
 }  // namespace communicator
 }  // namespace distributed
 }  // namespace raf
diff --git a/src/distributed/cuda/mpi_communicator.cc b/src/distributed/cuda/mpi_communicator.cc
@@ -62,6 +62,8 @@ MPICommunicator MPICommunicator::make(Value rank_list) {
 
 RAF_REGISTER_GLOBAL("raf.distributed.communicator._make.mpi").set_body_typed(MPICommunicator::make);
 
+RAF_REGISTER_OBJECT_REFLECT(MPICommunicatorObj);
+
 }  // namespace communicator
 }  // namespace distributed
 }  // namespace raf
diff --git a/src/distributed/cuda/nccl_communicator.cc b/src/distributed/cuda/nccl_communicator.cc
@@ -79,6 +79,8 @@ NCCLCommunicator NCCLCommunicator::make(Value rank_list) {
 RAF_REGISTER_GLOBAL("raf.distributed.communicator._make.nccl")
     .set_body_typed(NCCLCommunicator::make);
 
+RAF_REGISTER_OBJECT_REFLECT(NCCLCommunicatorObj);
+
 }  // namespace communicator
 }  // namespace distributed
 }  // namespace raf