From 4c6149236231ad6c7b2425946f94a16a8e71ae65 Mon Sep 17 00:00:00 2001
From: Yesir <135974443+YesianRohn@users.noreply.github.com>
Date: Wed, 29 May 2024 18:36:42 +0800
Subject: [PATCH] Fix reshape dimensions from (h, h) to (h, w) for correct
 image reshaping

---
 models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models.py b/models.py
index c90eeba7..d5056593 100644
--- a/models.py
+++ b/models.py
@@ -227,7 +227,7 @@ def unpatchify(self, x):
 
         x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
         x = torch.einsum('nhwpqc->nchpwq', x)
-        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
         return imgs
 
     def forward(self, x, t, y):