diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
index aa300b7210ed..0f95a6641912 100644
--- a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
+++ b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
@@ -193,6 +193,9 @@ public sealed class TypeEvalContext {
 
   @ApiStatus.Internal
   public <R> @Nullable R assumeType(@NotNull PyTypedElement element, @Nullable PyType type, @NotNull Function<TypeEvalContext, R> func) {
+    if (!Registry.is("python.use.better.control.flow.type.inference")) {
+      return func.apply(this);
+    }
     if (getKnownType(element) != null) {
       // Temporary solution, as overwriting known type might introduce inconsistencies with its dependencies.
       return null;