From 19f7cf620d33a107ca13cd076eea6a7cb1283b56 Mon Sep 17 00:00:00 2001
From: Sefik Ilkin Serengil <serengil@gmail.com>
Date: Mon, 18 Dec 2023 11:56:56 +0000
Subject: [PATCH] Update DeepFace.py

build_model return type
---
 deepface/DeepFace.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/deepface/DeepFace.py b/deepface/DeepFace.py
index d47008bef..337dad998 100644
--- a/deepface/DeepFace.py
+++ b/deepface/DeepFace.py
@@ -43,10 +43,13 @@
 tf_version = int(tf.__version__.split(".", maxsplit=1)[0])
 if tf_version == 2:
     tf.get_logger().setLevel(logging.ERROR)
+    from tensorflow.keras.models import Model
+else:
+    from keras.models import Model
 # -----------------------------------
 
 
-def build_model(model_name: str) -> Any:
+def build_model(model_name: str) -> Union[Model, Any]:
     """
     This function builds a deepface model
     Parameters:
@@ -55,7 +58,7 @@ def build_model(model_name: str) -> Any:
                     Age, Gender, Emotion, Race for facial attributes
 
     Returns:
-            built deepface model (keras.engine.Functional)
+            built deepface model ( (tf.)keras.models.Model )
     """
 
     # singleton design pattern