add conditional placement of feature column on cpu

alibaba · Feb 26, 2024 · 2955315 · 2955315
1 parent 93a25a6
commit 2955315
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/easy_rec/python/compat/feature_column/feature_column.py b/easy_rec/python/compat/feature_column/feature_column.py
@@ -171,6 +171,7 @@
 from tensorflow.python.util import nest
 
 from easy_rec.python.compat.feature_column import utils as fc_utils
+from easy_rec.python.utils import conditional
 from easy_rec.python.utils import constant
 from easy_rec.python.utils import embedding_utils
 
@@ -634,7 +635,7 @@ def _get_var_type(column):
       if embedding_utils.is_embedding_parallel():
         return _get_logits_embedding_parallel()
       else:
-        with ops.device('/CPU:0'):
+        with conditional(embedding_utils.embedding_on_cpu(), '/cpu:0'):
           return _get_logits()