Merge pull request #3 from gty111/batch

Fix batch dimension

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)