From a3190f334f3617d2032c470eb00eb01cc2bb19b7 Mon Sep 17 00:00:00 2001
From: wugunjun <985091524@qq.com>
Date: Tue, 16 May 2023 06:53:07 +0800
Subject: [PATCH] fix-total_variation_kernel_bug

---
 lib/cuda/total_variation_kernel.cu | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/cuda/total_variation_kernel.cu b/lib/cuda/total_variation_kernel.cu
index 72fab09..7a49795 100644
--- a/lib/cuda/total_variation_kernel.cu
+++ b/lib/cuda/total_variation_kernel.cu
@@ -28,8 +28,8 @@ __global__ void total_variation_add_grad_cuda_kernel(
     grad_to_add += (k==sz_k-1 ? 0 : wz * clamp(param[index]-param[index+1], -1.f, 1.f));
     grad_to_add += (j==0      ? 0 : wy * clamp(param[index]-param[index-sz_k], -1.f, 1.f));
     grad_to_add += (j==sz_j-1 ? 0 : wy * clamp(param[index]-param[index+sz_k], -1.f, 1.f));
-    grad_to_add += (i==0      ? 0 : wz * clamp(param[index]-param[index-sz_k*sz_j], -1.f, 1.f));
-    grad_to_add += (i==sz_i-1 ? 0 : wz * clamp(param[index]-param[index+sz_k*sz_j], -1.f, 1.f));
+    grad_to_add += (i==0      ? 0 : wx * clamp(param[index]-param[index-sz_k*sz_j], -1.f, 1.f));
+    grad_to_add += (i==sz_i-1 ? 0 : wx * clamp(param[index]-param[index+sz_k*sz_j], -1.f, 1.f));
     grad[index] += grad_to_add;
   }
 }