diff --git a/arch/csky/kernel/probes/uprobes.c b/arch/csky/kernel/probes/uprobes.c
index 2d31a12e46cfee0bcb202957a4c4b8eb3bcca223..936bea6fd32d191d349a288c8874db0eebc7c4bf 100644
--- a/arch/csky/kernel/probes/uprobes.c
+++ b/arch/csky/kernel/probes/uprobes.c
@@ -64,6 +64,7 @@ int arch_uprobe_post_xol(struct arch_uprobe *auprobe, struct pt_regs *regs)
 	struct uprobe_task *utask = current->utask;
 
 	WARN_ON_ONCE(current->thread.trap_no != UPROBE_TRAP_NR);
+	current->thread.trap_no = utask->autask.saved_trap_no;
 
 	instruction_pointer_set(regs, utask->vaddr + auprobe->insn_size);
 
@@ -101,6 +102,8 @@ void arch_uprobe_abort_xol(struct arch_uprobe *auprobe, struct pt_regs *regs)
 {
 	struct uprobe_task *utask = current->utask;
 
+	current->thread.trap_no = utask->autask.saved_trap_no;
+
 	/*
 	 * Task has received a fatal signal, so reset back to probed
 	 * address.