Problem with multiple gpu inference

Dear author.

Great work, but I encountered a few issues when I try to do multiple gpu inference. The code works well with 3b 7b inference on my single A800, but it got freezes during the encode part when I do 4 A800 inference with 1080p videos. Mutiple gpu works well with 720p videos, but it just got freezed during encode when it goes to 4 A800 infer 1080p videos. I think this could be related to the all_gather_base operations you used when gathering data, I tried to change them back to all_gather_tensor, but it still get freezes when I run with sp4 4 A800GPUS. Any idea how can I change this?

warnings.warn("The pts_unit 'pts' gives wrong results. Please use pts_unit 'sec'.")
Read video size: torch.Size([81, 3, 1136, 800])
Read video size: torch.Size([81, 3, 1136, 800])
Read video size: torch.Size([81, 3, 1136, 800])
Read video size: torch.Size([81, 3, 1136, 800])
Encoding videos: [torch.Size([3, 81, 1712, 1200])]
Encoding videos: [torch.Size([3, 81, 1712, 1200])]
Encoding videos: [torch.Size([3, 81, 1712, 1200])]
Encoding videos: [torch.Size([3, 81, 1712, 1200])]