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

chore(example/neural-style-transfer): GPU device for colab test

Browse Source
This commit is contained in:
sugarme 2020-07-05 01:44:34 +10:00
parent 8fcda3e1f9
commit 1eb4f57f6f
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
example/neural-style-transfer/main.go
View File

@ -86,10 +86,10 @@ func main() {
log.Fatal(err)
}
// cuda := gotch.CudaBuilder(0)
// device := cuda.CudaIfAvailable()
cuda := gotch.CudaBuilder(0)
device := cuda.CudaIfAvailable()
device := gotch.CPU
// device := gotch.CPU
netVS := nn.NewVarStore(device)
in := vision.NewImageNet()
net := vision.VGG16(netVS.Root(), in.ClassCount())