From 1a10fa0c82d5a7588bc168996df2f0e2fc40c726 Mon Sep 17 00:00:00 2001
From: Yao Matrix <matrix.yao@intel.com>
Date: Mon, 19 May 2025 08:35:32 +0800
Subject: [PATCH] enhance value guard of _device_agnostic_dispatch (#11553)

enhance value guard

Signed-off-by: Matrix Yao <matrix.yao@intel.com>
---
 src/diffusers/utils/testing_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py
index 00aad9d71a..ec220ff973 100644
--- a/src/diffusers/utils/testing_utils.py
+++ b/src/diffusers/utils/testing_utils.py
@@ -1212,8 +1212,8 @@ def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable],
 
     # Some device agnostic functions return values. Need to guard against 'None' instead at
     # user level
-    if fn is None:
-        return None
+    if not callable(fn):
+        return fn
 
     return fn(*args, **kwargs)