update for release

kailaix · Apr 28, 2020 · 2e95f27 · 2e95f27 · kailaix · Apr 28, 2020
1 parent 45503b6
commit 2e95f27
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/examples/gpu_custom_op/gputest.jl b/examples/gpu_custom_op/gputest.jl
@@ -1,9 +1,4 @@
 using ADCME
-using PyCall
-using LinearAlgebra
-using PyPlot
-using Random
-Random.seed!(233)
 
 function gpu_test(a)
  gpu_test_ = load_op_and_grad("$(@__DIR__)/build/libGpuTest","gpu_test")
@@ -12,7 +7,7 @@ function gpu_test(a)
 end
 
 # TODO: specify your input parameters
-a = rand(100)
+a = rand(3)
 u = gpu_test(a)
 sess = Session(); init(sess)
 v1 = run(sess, u)

diff --git a/src/extra.jl b/src/extra.jl
@@ -732,7 +732,14 @@ and `nvcc` is in your path.""")
 
 end
 
+
+"""
+ test_gpu()
+
+Tests the GPU ultilities
+"""
 function test_gpu()
+ PWD = pwd()
  run(`which nvcc`)
  cd("$(@__DIR__)/../examples/gpu_custom_op")
  mkdir("build")
@@ -741,4 +748,5 @@ function test_gpu()
  ADCME.make()
  cd("..")
  include("gputest.jl")
+ cd(PWD)
 end