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

updated

Browse Source
This commit is contained in:
sugarme 2020-07-08 20:53:58 +10:00
parent 11d345360d
commit f76566834c
2 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
example/cifar/main.go
View File

@ -107,9 +107,9 @@ func main() {
startRAM := si.TotalRam - si.FreeRam
// cuda := gotch.CudaBuilder(0)
// device := cuda.CudaIfAvailable()
device := gotch.CPU
cuda := gotch.CudaBuilder(0)
device := cuda.CudaIfAvailable()
// device := gotch.CPU
vs := nn.NewVarStore(device)

1
tensor/data.go
View File

@ -101,6 +101,7 @@ func (it *Iter2) Shuffle() {
it.xs = it.xs.MustIndexSelect(0, index, true)
it.ys = it.ys.MustIndexSelect(0, index, true)
index.MustDrop()
}
// ToDevice transfers the mini-batches to a specified device.