[Bugfix] Minor fix in builder.py (#1235)

Author: LJC00118

tilelang/language/v2/builder.py

@@ -390,6 +390,7 @@ def aug_assign(self, op, target, aug_value):
         self.check_continue_break()
         if is_var(target):
             tir.buffer_store(target, eval_op(op, target[0], aug_value), 0)
+            return target
         elif isinstance(target, Buffer):
             raise RuntimeError("Augmented assignment is not supported for Buffer")
         else: