Fix batch dimension

xdit-project · Sep 23, 2024 · a067f21 · a067f21
1 parent 90a4e0b
commit a067f21
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/distvae/modules/patch_utils.py b/distvae/modules/patch_utils.py
@@ -34,7 +34,7 @@ def forward(self, patch_hidden_state):
         )
         patch_hidden_state_list = [
             torch.empty(
-                [1, patch_hidden_state.shape[1], patch_height_list[i].item(), patch_hidden_state.shape[-1]], 
+                [patch_hidden_state.shape[0], patch_hidden_state.shape[1], patch_height_list[i].item(), patch_hidden_state.shape[-1]], 
                 dtype=patch_hidden_state.dtype,
                 device=f"cuda:{self.rank}"
             ) for i in range(self.world_size)