fix bug of negative sampling run in ps mode with just 2 workers (#501)

* fix bug of negative sampling run in ps mode with just 2 workers
alibaba · Nov 13, 2024 · 94dce1a · 94dce1a
1 parent bad99f5
commit 94dce1a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/easy_rec/python/core/sampler.py b/easy_rec/python/core/sampler.py
@@ -79,7 +79,10 @@ def _init_graph(self):
       if 'ps' in tf_config['cluster']:
         # ps mode
         tf_config = json.loads(os.environ['TF_CONFIG'])
-        task_count = len(tf_config['cluster']['worker']) + 2
+        if 'worker' in tf_config['cluster']:
+          task_count = len(tf_config['cluster']['worker']) + 2
+        else:
+          task_count = 2
         if self._is_on_ds:
           gl.set_tracker_mode(0)
           server_hosts = [

diff --git a/requirements/docs.txt b/requirements/docs.txt
@@ -5,4 +5,4 @@ recommonmark==0.6.0
 sphinx==5.1.1
 sphinx_markdown_tables==0.0.17
 sphinx_rtd_theme
-tensorflow-probability
+tensorflow-probability==0.11.0