浏览代码

add some docs for remoteexpert forward

justheuristic 5 年之前
父节点
当前提交
e050714f56
共有 2 个文件被更改,包括 2 次插入0 次删除
  1. 1 0
      docs/modules/client.rst
  2. 1 0
      tesseract/client/remote_expert.py

+ 1 - 0
docs/modules/client.rst

@@ -12,6 +12,7 @@ tesseract.client
   <br><br>
 
 .. autoclass:: RemoteExpert
+   :members: forward
 
 .. autoclass:: GatingFunction
    :members:

+ 1 - 0
tesseract/client/remote_expert.py

@@ -25,6 +25,7 @@ class RemoteExpert(nn.Module):
         self._info = None
 
     def forward(self, *args, **kwargs):
+        """ Call RemoteExpert for the specified inputs and return its output(s). Compatible with pytorch.autograd. """
         assert len(kwargs) == len(self.info['keyword_names']), f"Keyword args should be {self.info['keyword_names']}"
         kwargs = {key: kwargs[key] for key in self.info['keyword_names']}
         # Note: we put keyword arguments in the same order as on a server to prevent f(a=1, b=2) != f(b=2, a=1) errors