diff --git a/example/augmentation/README.md b/example/augmentation/README.md
index f20d32e..825d6df 100644
--- a/example/augmentation/README.md
+++ b/example/augmentation/README.md
@@ -28,4 +28,7 @@ There are 2 APIs (`aug.Compose` and `aug.OneOf`) to compose augmentation methods
 ```
 
 
+![transformed images](augment.png)
+
+
 
diff --git a/example/augmentation/augment.png b/example/augmentation/augment.png
new file mode 100644
index 0000000..d047cbc
Binary files /dev/null and b/example/augmentation/augment.png differ
diff --git a/example/augmentation/main.go b/example/augmentation/main.go
index c50ec1d..00f5c60 100644
--- a/example/augmentation/main.go
+++ b/example/augmentation/main.go
@@ -16,8 +16,8 @@ func main() {
 			panic(err)
 		}
 
-		device := gotch.CudaIfAvailable()
-		// device := gotch.CPU
+		// device := gotch.CudaIfAvailable()
+		device := gotch.CPU
 		imgTs := img.MustTo(device, true)
 		// t, err := aug.Compose(aug.WithResize(512, 512)) // NOTE. WithResize just works on CPU.
 		// t, err := aug.Compose(aug.WithRandRotate(0, 360), aug.WithColorJitter(0.3, 0.3, 0.3, 0.4))
@@ -37,6 +37,7 @@ func main() {
 		// t, err := aug.Compose(aug.WithNormalize(aug.WithNormalizeMean([]float64{0.485, 0.456, 0.406}), aug.WithNormalizeStd([]float64{0.229, 0.224, 0.225})))
 
 		t, err := aug.Compose(
+			aug.WithResize(200, 200),
 			aug.WithRandomVFlip(0.5),
 			aug.WithRandomHFlip(0.5),
 			aug.WithRandomCutout(),
diff --git a/vision/aug/function.go b/vision/aug/function.go
index 8464bd9..250bfa9 100644
--- a/vision/aug/function.go
+++ b/vision/aug/function.go
@@ -761,7 +761,6 @@ func applyGridTransform(x, gridInput *ts.Tensor, mode string, fillValue []float6
 	addTs := mask.MustMul1(ts.FloatScalar(-1), false).MustAdd1(ts.FloatScalar(1.0), true).MustMul(fillImg, true)
 	imgOut := image.MustMul(mask, true).MustAdd(addTs, true)
 	addTs.MustDrop()
-	image.MustDrop()
 	mask.MustDrop()
 	fillImg.MustDrop()