diff --git a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponedArgumentsAnalyzer.kt b/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponedArgumentsAnalyzer.kt
index 5e07434ce68..9ee5b599b04 100644
--- a/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponedArgumentsAnalyzer.kt
+++ b/compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponedArgumentsAnalyzer.kt
@@ -98,7 +98,9 @@ class PostponedArgumentsAnalyzer(
 
         if (returnArguments.isEmpty()) {
             val unitType = lambda.returnType.builtIns.unitType
-            c.getBuilder().addSubtypeConstraint(lambda.returnType.let(::substitute), unitType, LambdaArgumentConstraintPosition(lambda))
+            val lambdaReturnType = lambda.returnType.let(::substitute)
+            c.getBuilder().addSubtypeConstraint(lambdaReturnType, unitType, LambdaArgumentConstraintPosition(lambda))
+            c.getBuilder().addSubtypeConstraint(unitType, lambdaReturnType, LambdaArgumentConstraintPosition(lambda))
         }
 
         lambda.setAnalyzedResults(returnArguments, subResolvedKtPrimitives)
diff --git a/compiler/testData/diagnostics/tests/inference/kt28654.kt b/compiler/testData/diagnostics/tests/inference/kt28654.kt
index 5372d7e356a..38096364d53 100644
--- a/compiler/testData/diagnostics/tests/inference/kt28654.kt
+++ b/compiler/testData/diagnostics/tests/inference/kt28654.kt
@@ -2,7 +2,7 @@
 // !DIAGNOSTICS: -UNUSED_VARIABLE
 // Related issue: KT-28654
 
-fun <K> select(): K = <!OI;TYPE_INFERENCE_EXPECTED_TYPE_MISMATCH!>run { <!OI;TYPE_MISMATCH!><!>}<!>
+fun <K> select(): K = <!NI;TYPE_MISMATCH, OI;TYPE_INFERENCE_EXPECTED_TYPE_MISMATCH!>run { <!OI;TYPE_MISMATCH!><!>}<!>
 
 fun test() {
     val x: Int = select()