fix(example/neural-style-transfer): save image

2020-07-05 15:02:45 +10:00 · 2020-07-05 15:02:45 +10:00 · a0653fe4d3
commit a0653fe4d3
parent 8fc6efca1c
1 changed files with 1 additions and 1 deletions
--- a/example/neural-style-transfer/main.go
+++ b/example/neural-style-transfer/main.go
@ -175,7 +175,7 @@ func main() {
 		if (stepIdx % 1000) == 0 {
 			clone := inputVar.MustShallowClone()
 			img := clone.MustDetach()
-			imageTs := img.MustTo(device, true)
+			imageTs := img.MustTo(gotch.CPU, true)
 			clone.MustDrop()
 			err := in.SaveImage(imageTs, fmt.Sprintf("../../data/neural-style-transfer/out%v.jpg", stepIdx))
 			if err != nil {