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