fixed bugs

justanhduc · Jun 17, 2021 · cfb5d01 · cfb5d01
1 parent 2bff290
commit cfb5d01
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/neural_monitor/monitor.py b/neural_monitor/monitor.py
@@ -208,7 +208,7 @@ def func(self, name: str, *args, **kwargs):
 
 def distributed_collect(f):
     @functools.wraps(f)
-    def func(self: Monitor, name: str, value: T.Tensor, *args, **kwargs):
+    def func(self, name: str, value: T.Tensor, *args, **kwargs):
         if self.distributed:
             assert isinstance(value, T.Tensor), 'value must be a Tensor in distributed mode'
 
@@ -221,7 +221,7 @@ def func(self: Monitor, name: str, value: T.Tensor, *args, **kwargs):
 
 def distributed_flush(f):
     @functools.wraps(f)
-    def func(self: Monitor, *args, **kwargs):
+    def func(self, *args, **kwargs):
         if self.distributed and self.rank != 0:
             return
         return f(*args, **kwargs)