fix ensure_bagua_tensor

BaguaSys · Oct 21, 2021 · 2aac918 · 2aac918
1 parent dfd81f9
commit 2aac918
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/bagua/torch_api/distributed.py b/bagua/torch_api/distributed.py
@@ -294,9 +294,7 @@ def with_bagua(  # pytype: disable=module-attr
             ...      GradientAllReduce()
             ...    )
         """
-        self.bagua_module_name = "{}_{}".format(
-            self.__class__.__name__, next(BaguaModule.__id_iter)
-        )
+        self.bagua_module_name = "{}_{}".format(self.__class__.__name__, id(self))
 
         # set bucket process group
         if process_group is None:

diff --git a/bagua/torch_api/tensor.py b/bagua/torch_api/tensor.py
@@ -43,6 +43,11 @@ def ensure_bagua_tensor(
                     self.bagua_tensor_name == name
                 ), "assigning a different name to existing bagua tensor is forbidden"
 
+                assert (
+                    self.bagua_module_name == module_name
+                ), "assigning a different module name to existing bagua tensor is forbidden"
+            return self
+
         self.bagua_tensor_name = name if name is not None else ""
         self.bagua_module_name = module_name
         self.bagua_backend = (