andr3/gotch
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fixed initializing CUDA device at example/mnist cnn model

Browse Source
This commit is contained in:
sugarme 2021-07-23 16:26:39 +10:00
parent 7e9eb5bff2
commit 8baf32d860
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
example/mnist/cnn.go
View File

@ -74,9 +74,8 @@ func runCNN1() {
testImages := ds.TestImages testImages := ds.TestImages
testLabels := ds.TestLabels testLabels := ds.TestLabels
cuda := gotch.CudaBuilder(0) device := gotch.CudaIfAvailable()
vs := nn.NewVarStore(cuda.CudaIfAvailable()) vs := nn.NewVarStore(device)
// vs := nn.NewVarStore(gotch.CPU)
net := newNet(vs.Root()) net := newNet(vs.Root())
opt, err := nn.DefaultAdamConfig().Build(vs, LrCNN) opt, err := nn.DefaultAdamConfig().Build(vs, LrCNN)