fix exception segmentfault

mindspore/core/ir/tensor.h

@@ -87,10 +87,9 @@ class WaitEvent : public ExceptionListener {
     if (!need_wait_) {
       return;
     }
-    MsException::Instance().AddExceptionListener(const_cast<WaitEvent *>(this));
+    MsException::Instance().SetExceptionListener(const_cast<WaitEvent *>(this));
     cond_var_.wait(lock, [this] { return !need_wait_; });
     MsException::Instance().CheckException();
-    MsException::Instance().RemoveExceptionListener(const_cast<WaitEvent *>(this));
   }
 
   void set_need_wait(bool need_wait) {

mindspore/core/utils/ms_exception.h

@@ -34,13 +34,10 @@ class MsException {
 
   void SetException() {
     exception_ptr_ = std::current_exception();
-    if (exception_ptr_ != nullptr) {
-      for (auto &listener : listeners_) {
-        if (listener == nullptr) {
-          continue;
-        }
-        listener->OnException();
-      }
+    if (exception_ptr_ != nullptr && listener_ != nullptr) {
+      auto listener = listener_;
+      listener_ = nullptr;
+      listener->OnException();
     }
   }
 
@@ -52,15 +49,13 @@ class MsException {
     }
   }
 
-  void AddExceptionListener(ExceptionListener *listener) { (void)listeners_.insert(listener); }
-
-  void RemoveExceptionListener(ExceptionListener *listener) { (void)listeners_.erase(listener); }
+  void SetExceptionListener(ExceptionListener *listener) { listener_ = listener; }
 
  private:
   MsException() = default;
   ~MsException() = default;
   DISABLE_COPY_AND_ASSIGN(MsException)
-  std::set<ExceptionListener *> listeners_;
+  ExceptionListener *listener_{nullptr};
   std::exception_ptr exception_ptr_{nullptr};
 };
 }  // namespace mindspore