Add Span to Config

variable/config.go

@@ -10,16 +10,24 @@ var Config = config{
 	Train:          true,
 }
 
-func Nograd() func() {
+type Span struct {
+	End func()
+}
+
+func Nograd() *Span {
 	Config.EnableBackprop = false
-	return func() {
-		Config.EnableBackprop = true
+	return &Span{
+		End: func() {
+			Config.EnableBackprop = true
+		},
 	}
 }
 
-func TestMode() func() {
+func TestMode() *Span {
 	Config.Train = false
-	return func() {
-		Config.Train = true
+	return &Span{
+		End: func() {
+			Config.Train = true
+		},
 	}
 }

variable/config_test.go

@@ -20,7 +20,7 @@ func ExampleNograd() {
 	f()
 
 	func() {
-		defer variable.Nograd()()
+		defer variable.Nograd().End()
 
 		fmt.Println("backprop:", variable.Config.EnableBackprop)
 		f()
@@ -44,7 +44,7 @@ func ExampleTestMode() {
 	fmt.Println("train:", variable.Config.Train)
 
 	func() {
-		defer variable.TestMode()()
+		defer variable.TestMode().End()
 
 		fmt.Println("train:", variable.Config.Train)
 	}()

variable/variable.go

@@ -63,7 +63,7 @@ func (v *Variable) Backward(opts ...Opts) {
 		// backward
 		func() {
 			if len(opts) == 0 || !opts[0].CreateGraph {
-				defer Nograd()()
+				defer Nograd().End()
 			}
 
 			gxs := f.Backward(gys...)