From 8031bb85f5943fb587d27ca66c0caf4809e12765 Mon Sep 17 00:00:00 2001
From: pjongy <hi.pjongy@gmail.com>
Date: Thu, 19 Oct 2023 03:31:22 +0000
Subject: [PATCH] MINOR: Move runtime.SetFinalizer into freeXXX

---
 ts/scalar.go | 12 ++++++++----
 ts/tensor.go | 14 ++++++++------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/ts/scalar.go b/ts/scalar.go
index 33da8c9..c548637 100644
--- a/ts/scalar.go
+++ b/ts/scalar.go
@@ -17,8 +17,10 @@ type Scalar struct {
 
 // free releases C allocated memory.
 func freeCScalar(x *Scalar) error {
-	nbytes := x.nbytes()
-	atomic.AddInt64(&AllocatedMem, -nbytes)
+	if gotch.Debug {
+		nbytes := x.nbytes()
+		atomic.AddInt64(&AllocatedMem, -nbytes)
+	}
 	lock.Lock()
 	delete(ExistingScalars, x.name)
 	lock.Unlock()
@@ -53,8 +55,10 @@ func newScalar(cscalar lib.Cscalar, nameOpt ...string) *Scalar {
 	}
 
 	atomic.AddInt64(&ScalarCount, 1)
-	nbytes := x.nbytes()
-	atomic.AddInt64(&AllocatedMem, nbytes)
+	if gotch.Debug {
+		nbytes := x.nbytes()
+		atomic.AddInt64(&AllocatedMem, nbytes)
+	}
 	lock.Lock()
 	ExistingScalars[x.name] = struct{}{}
 	lock.Unlock()
diff --git a/ts/tensor.go b/ts/tensor.go
index 33f6666..dc8e554 100644
--- a/ts/tensor.go
+++ b/ts/tensor.go
@@ -67,8 +67,10 @@ func newTensor(ctensor lib.Ctensor, nameOpt ...string) *Tensor {
 	x.d = new(bigStruct)
 
 	atomic.AddInt64(&TensorCount, 1)
-	nbytes := x.nbytes()
-	atomic.AddInt64(&AllocatedMem, nbytes)
+	if gotch.Debug {
+		nbytes := x.nbytes()
+		atomic.AddInt64(&AllocatedMem, nbytes)
+	}
 	lock.Lock()
 	if _, ok := ExistingTensors[name]; ok {
 		name = fmt.Sprintf("%s_%09d", name, TensorCount)
@@ -170,6 +172,10 @@ func freeCTensor(ts *Tensor) error {
 	// IMPORTANT. make it nil so won't double free.
 	ts.ctensor = nil
 
+	// Clear SetFinalizer on ts so no double free tensor.
+	// Ref. https://pkg.go.dev/runtime#SetFinalizer
+	runtime.SetFinalizer(ts, nil)
+
 	return nil
 }
 
@@ -1190,10 +1196,6 @@ func (ts *Tensor) Drop() error {
 		return nil
 	}
 
-	// Clear SetFinalizer on ts so no double free tensor.
-	// Ref. https://pkg.go.dev/runtime#SetFinalizer
-	runtime.SetFinalizer(ts, nil)
-
 	ts.calledFrom = "ts.Drop()"
 	return freeCTensor(ts)
 }