diff --git a/drivers/gpu/drm/i915/selftests/igt_flush_test.c b/drivers/gpu/drm/i915/selftests/igt_flush_test.c
index 7f35bddc2e9522ca8cb70a80374df8dc65f4baf9..0d06f559243f9e7c462fe557a0fdd7ce11f88240 100644
--- a/drivers/gpu/drm/i915/selftests/igt_flush_test.c
+++ b/drivers/gpu/drm/i915/selftests/igt_flush_test.c
@@ -57,7 +57,8 @@ int igt_flush_test(struct drm_i915_private *i915, unsigned int flags)
 
 	cond_resched();
 
-	if (i915_gem_switch_to_kernel_context(i915)) {
+	if (flags & I915_WAIT_LOCKED &&
+	    i915_gem_switch_to_kernel_context(i915)) {
 		pr_err("Failed to switch back to kernel context; declaring wedged\n");
 		i915_gem_set_wedged(i915);
 	}