cleanup, fix warnings, doc

rwth-i6 · Nov 28, 2023 · 88d53f3 · 88d53f3
1 parent 90c7548
commit 88d53f3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/returnn/torch/distributed.py b/returnn/torch/distributed.py
@@ -8,13 +8,9 @@
 import os
 import socket
 
-from contextlib import contextmanager
 import torch
-from torch.distributed.algorithms.join import Join
-from torch.nn.parallel import DistributedDataParallel
 
 from returnn.config import Config
-import returnn.frontend as rf
 
 
 class DistributedContext:

diff --git a/returnn/torch/engine.py b/returnn/torch/engine.py
@@ -1057,6 +1057,11 @@ def __init__(self, *, module: torch.nn.Module, engine: Engine):
         self.engine = engine
 
     def forward(self, *args, **kwargs):
+        """
+        Call run step function (train_step or forward_step).
+
+        :return: all produced raw tensors via the run context (:func:`rf.get_run_ctx`).
+        """
         # noinspection PyProtectedMember
         self.engine._run_step(*args, **kwargs, _inside_wrapped=True)