From e71258f19f228b8825d7e42ad6e96a321a0e255c Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Wed, 3 Jul 2024 02:59:44 +0800
Subject: [PATCH] 'update'

---
 hook_HYDiT_utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py
index a37f9d4..a2d9d25 100644
--- a/hook_HYDiT_utils.py
+++ b/hook_HYDiT_utils.py
@@ -333,7 +333,10 @@ class CustomizeEmbedsModel(nn.Module):
     def forward(self, *args, **kwargs):
         input_ids = kwargs.get("input_ids", None)
         if self.x is None:
-            batch_size = input_ids.shape[0]
+            if input_ids is None:
+                batch_size = 1
+            else:
+                batch_size = input_ids.shape[0]
             self.x = torch.zeros(1, batch_size, 256, 2048, dtype=self.dtype)
 
         if kwargs.get("output_hidden_states", False):