diff --git a/drivers/misc/lkdtm.c b/drivers/misc/lkdtm.c
index 28adefe..5cbd740 100644
--- a/drivers/misc/lkdtm.c
+++ b/drivers/misc/lkdtm.c
@@ -315,7 +315,7 @@ static void lkdtm_do_action(enum ctype which)
 	}
 	case CT_OVERWRITE_ALLOCATION: {
 		size_t len = 1020;
-		u32 *data = kmalloc(len, GFP_KERNEL);
+		u32 *data = kmalloc(len, GFP_ATOMIC);
 
 		data[1024 / sizeof(u32)] = 0x12345678;
 		kfree(data);
@@ -323,14 +323,16 @@ static void lkdtm_do_action(enum ctype which)
 	}
 	case CT_WRITE_AFTER_FREE: {
 		size_t len = 1024;
-		u32 *data = kmalloc(len, GFP_KERNEL);
+		u32 *data = kmalloc(len, GFP_ATOMIC);
 
 		kfree(data);
-		schedule();
+		udelay(100);
 		memset(data, 0x78, len);
 		break;
 	}
 	case CT_SOFTLOCKUP:
+		if (in_interrupt())
+			break;
 		preempt_disable();
 		for (;;)
 			cpu_relax();
@@ -341,6 +343,8 @@ static void lkdtm_do_action(enum ctype which)
 			cpu_relax();
 		break;
 	case CT_HUNG_TASK:
+		if (in_interrupt())
+			break;
 		set_current_state(TASK_UNINTERRUPTIBLE);
 		schedule();
 		break;