From 0a398b12196e8b83b6065a9a1b21988d600c46fc Mon Sep 17 00:00:00 2001
From: LielinJiang <jianglielin@baidu.com>
Date: Tue, 7 Apr 2020 14:34:12 +0000
Subject: [PATCH] add docs for set_device

---
 model.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/model.py b/model.py
index e6faeb7..6fecbf1 100644
--- a/model.py
+++ b/model.py
@@ -42,6 +42,14 @@ __all__ = ['Model', 'Loss', 'CrossEntropy', 'Input', 'set_device']
 
 
 def set_device(device):
+    """
+    Args:
+        device (str): specify device type, 'cpu' or 'gpu'.
+        
+    Returns:
+        fluid.CUDAPlace or fluid.CPUPlace: Created GPU or CPU place.
+    """
+
     assert isinstance(device, six.string_types) and device.lower() in ['cpu', 'gpu'], \
     "Expected device in ['cpu', 'gpu'], but got {}".format(device)
 
@@ -1082,7 +1090,11 @@ class Model(fluid.dygraph.Layer):
 
         return eval_result
 
-    def predict(self, test_data, batch_size=1, num_workers=0, stack_outputs=True):
+    def predict(self,
+                test_data,
+                batch_size=1,
+                num_workers=0,
+                stack_outputs=True):
         """
         FIXME: add more comments and usage
         Args:
-- 
GitLab