From 6983485eed90d631c593c634f0bcee9d780b089d Mon Sep 17 00:00:00 2001
From: sayakpaul <spsayakpaul@gmail.com>
Date: Mon, 12 Jan 2026 13:47:32 +0530
Subject: [PATCH] up

---
 src/diffusers/training_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusers/training_utils.py b/src/diffusers/training_utils.py
index e96e0c5711..a36455a43f 100644
--- a/src/diffusers/training_utils.py
+++ b/src/diffusers/training_utils.py
@@ -432,7 +432,7 @@ def get_fsdp_kwargs_from_accelerator(accelerator) -> dict:
 
 def wrap_with_fsdp(
     model: torch.nn.Module,
-    device: str | torch.Device,
+    device: str | torch.device,
     offload: bool = True,
     use_orig_params: bool = True,
     limit_all_gathers: bool = True,