nvidia_test

TransBiz · May 12, 2017 · 5acedf5 · 5acedf5
1 parent 25feba9
commit 5acedf5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -62,5 +62,7 @@ EXPOSE 6006
 EXPOSE 8888
 
 WORKDIR "/home"
+COPY nvidia_test.py 
+CMD python nvidia_test.py
 
 CMD ["/run_jupyter.sh", "--allow-root"]
diff --git a/nvidia_test.py b/nvidia_test.py
@@ -0,0 +1,10 @@
+import tensorflow as tf
+
+# Creates a graph.
+a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
+b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
+c = tf.matmul(a, b)
+# Creates a session with log_device_placement set to True.
+sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
+# Runs the op.
+print(sess.run(c))