fix type error

alibaba · Feb 27, 2024 · bb5e69e · bb5e69e
1 parent 2955315
commit bb5e69e
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
@@ -635,7 +635,8 @@ def _get_var_type(column):
       if embedding_utils.is_embedding_parallel():
         return _get_logits_embedding_parallel()
       else:
-        with conditional(embedding_utils.embedding_on_cpu(), '/cpu:0'):
+        with conditional(embedding_utils.embedding_on_cpu(),
+                         ops.device('/cpu:0')):
           return _get_logits()