How to generate on multiple GPUs?
Hi!

How would I run generation on multiple GPUs at the same time? Running model.generate on a DataParallel layer isn't possible, and model.module.generate run on a single GPU. 

Any advice would be appreciated!
Hey @moinnadeem,

Sorry to answer so late...yeah this will require some work I think! Will put it in the generate() project, but not sure when we manage to take a deeper look into this. Feel free to open a PR and tag me if you have some good ideas :-) <|||||>Just added this to `example/seq2seq/`: https://github.com/huggingface/transformers/blob/master/examples/seq2seq/README.md#multi-gpu-evalulation<|||||>+1 on getting `generate()` to run when using nn.dataparallel for evaluation

Is the current finetune.py script (which I think uses DDP) running generate on all GPUs or just on one device?<|||||>it uses all devices.<|||||>Regarding getting generate() to run using nn.dataparallel, it actually probably isn't worth building the functionality. It seems that dataparallel (e.g. with 2 GPUs) is rarely faster (and actually sometimes quite a bit slower) than running with a single GPU, even when the multi-gpu is on the same node.<|||||>Hello,

I want to use generate function with single GPU. Specifically, I fine tuned a GPT-2 model (on GPU) and subsequently, I want to generate text with it. 

When I run this 

```
input_ids.to(device)
sample_output = model.generate(
    input_ids, 
    do_sample=True, 
    max_length=150,
    top_k=50,
    top_p=0.92
)
```
I get this error

`RuntimeError: Input, output and indices must be on the current device`

When i move the model and the input `.to('cpu')`, it works.<|||||>@contribcode this looks like another issue, do you mind opening a new issue with the issue template filled-in? Thank you.<|||||>You are right @LysandreJik, I will open a new issue.<|||||>This issue has been stale for 1 month.<|||||>Gently pinging @stas00 here - is this possible now thanks to your recent PR in `generate()`? <|||||>1. The work I did in `generate`'s search functions is to make those work under deepspeed zero-3+ regime, where all gpus must work in sync to complete, even if some of them finished their sequence early - it uses all gpus because the params are sharded across all gpus and thus all gpus contribute their part to make it happen. 

2. But otherwise, the current design is that we already use all gpus. Is it not the case? We just don't do anything with the results from all but rank 0 process.

3. In current examples we for some reason save the generated tokens only from rank0: e.g. in `translation_run.py`:

https://github.com/huggingface/transformers/blob/50f4539b8201b26b18085260bf801cdeadfa6640/examples/pytorch/translation/run_translation.py#L564

I guess it's easier than to try to somehow write them out in non-interleaved way - probably could to add `flock` and append them all to the same file if needed.

4. metrics we calculate on all gpus, but only save metrics from rank0 process

https://github.com/huggingface/transformers/blob/5c00918681d6b4027701eb46cea8f795da0d4064/src/transformers/trainer_pt_utils.py#L924

If I missed something please clarify what is not working in `generate` under multiple gpus or what is the desired functionality?<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.<|||||>Any work on this? Is there any way to use DDP with generate?<|||||>Also curious about this<|||||>We should definitely enable `generate()` on multiple GPUs - did anybody give it a try with DDP? <|||||>See:  https://github.com/huggingface/transformers/issues/6821#issuecomment-825770020<|||||>> Hi!
> 
> How would I run generation on multiple GPUs at the same time? Running model.generate on a DataParallel layer isn't possible, and model.module.generate run on a single GPU.
> 
> Any advice would be appreciated!

Hello,
Did you resolve this issue? If you remember can you share me the solution for it?

Thanks in advance<|||||>As I already replied earlier `generate` works on multiple gpus including `deepspeed`<|||||>> 
@kkavyashankar0009 @JulesGM 

Hi, it seems that model.generatre() does not support DP, but it can be used in DDP. Here is some example code:


```
import torch.multiprocessing as mp
import torch.distributed as dist
import argparse
import torch
from transformers import get_linear_schedule_with_warmup, BartConfig, BartTokenizer, BartForConditionalGeneration

parser = argparse.ArgumentParser()
parser.add_argument('--nodes', type=int, default=1)  # how many nodes (machines) you have
parser.add_argument('--gpus', type=int, default=-1, help='num gpus per node')
parser.add_argument('--nr', type=int, default=0, help='ranking within the nodes')
args = parser.parse_args()
tokenizer = BartTokenizer.from_pretrained(args.tokenizer_name)

def test_model_generation(local_gpu_rank, args):
    set_seed(args.seed)
    args.rank = args.nr * args.gpus + local_gpu_rank  # compute the rank of the current GPU
    dist.init_process_group(backend="nccl", init_method="env://", world_size=args.world_size, rank=args.rank)

    test_data = "../data/" + args.dataset + "/test.txt"
    print("Processing data: " + test_data, flush=True)
    config = BartConfig.from_pretrained(args.config_name)
    bart_ctx = BartForConditionalGeneration.from_pretrained(args.model_name_or_path, config=config)
    bart_rep = BartForConditionalGeneration.from_pretrained(args.model_name_or_path, config=config)
    bart_ctx.resize_token_embeddings(len(tokenizer))
    bart_rep.resize_token_embeddings(len(tokenizer))

    model = MyModel(bart, config)
    model_state_dict = torch.load("../output/model/gen.ddp.pt")  
    model.load_state_dict(model_state_dict, strict=True)  # load model
    torch.cuda.set_device(local_gpu_rank)  
    args.device = torch.device("cuda", local_gpu_rank)
    model.to(args.device)  # move the model to GPU
    # model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_gpu_rank], find_unused_parameters=True)
    model.eval()

    test_dataset = FileDataset(test_data, tokenizer, max_context_len=args.max_ctx_len, max_response_len=args.max_rep_len, dataset=args.dataset)
    test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset, shuffle=False)
    test_dataloader = DataLoader(test_dataset, batch_size=args.test_batch_size, num_workers=2, sampler=test_sampler)  # the batch size on each GPU

    if args.rank == 0:
        count = 0
        fw = open("../output/test.responses.txt", "w", encoding="utf-8")  # open file for writting result
    with torch.no_grad():
        test_sampler.set_epoch(0)  # keep all data the same on all GPUs, it is usually used in training, I'm not sure if it is necessary in inference
        for test_data in test_dataloader:
            for key in test_data.keys():
                test_data[key] = test_data[key].to(args.device)
            outputs = model.bart_model.generate(
                input_ids=test_data["input_ids"],
                attention_mask=test_data["attention_mask"],
                max_length=args.max_rep_len,
                no_repeat_ngram_size=3,
                num_beams=10,
            )  # my model contains a BART model in self.bart_model，so I use model.module.bart_model to get it
            if outputs.size(1) < args.max_rep_len:  # need padding because the lengths from different GPUs may be different
                batch_pred_padding = torch.ones((outputs.size(0), args.max_rep_len - outputs.size(1)), dtype=outputs.dtype).cuda()  # use the padding token of BART, and its token id is 1. Be careful with the data type.
                outputs = torch.cat([outputs, batch_pred_padding], dim=1)
            batch_pred = [torch.zeros_like(outputs, dtype=outputs.dtype).cuda() for _ in range(args.world_size)]  # initialized a list for collecting tensors from all GPUs. Be careful with the data type.
            dist.all_gather(batch_pred, outputs)  # collect data
            batch_pred = torch.stack(batch_pred, dim=1)  # use stack, take care of the dimension
            batch_pred = batch_pred.reshape(-1, args.max_rep_len)
            if args.rank == 0:
                batch_out_sentences = tokenizer.batch_decode(batch_pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)  # decode the token id to token
                for r in batch_out_sentences:
                    fw.write(r + "\n")
                    fw.flush()
                count += len(batch_out_sentences)
                print(count)
    if args.rank == 0:
        fw.close()
        
if __name__ == "__main__":
    if args.gpus < 0:
        args.gpus = torch.cuda.device_count()
    args.world_size = args.nodes * args.gpus
    os.environ['MASTER_ADDR']='localhost'
    os.environ['MASTER_PORT']='8888'
    mp.spawn(test_model_generation, nprocs=args.gpus, args=(args, ))

```




Be careful with the data order. It is better to add IDs and write them to the file, which can be used for checking the order.<|||||>@DaoD 

Thanks a lot for this example!  If I understand correctly, I think you don't need to use `torch.nn.parallel.DistributedDataParallel` wrapper on the model if only running inference though, since we don't care about gradients.  I found that only using the distributed data sampler and leaving the model as is (without the DDP wrapper) led to better memory usage.  Haven't looked into why, but let's me use larger batch sizes just FYI.<|||||>@bkleiner2 Thanks for your information! I just want to use multiple GPUs for inference, so I don't how it works if DDP is not used. <|||||>Yeah I'm just saying for inference only I don't think you need this line:

`model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_gpu_rank], find_unused_parameters=True)`

because we don't need to aggregate gradients.  Once you've done this:

`    model.to(args.device)  # move the model to GPU` you've already copied the model to each device right?  Because there is a separate process running for each device.  So you can get the data parallelism across devices by simply using `test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset, shuffle=False)`, as far as I can tell.<|||||>@bkleiner2 Oh, I got it! Thanks for your explanation!<|||||>I think the real problem we are facing is this error when we try to called `model.generate`
``` bash
AttributeError: 'DistributedDataParallel' object has no attribute 'generate'
```

Indeed, it can be solved by distributed sampler or using the huggingface trainer class. 

But it could be quite tricky if we don't use them and write our own trainer. Even we have to be careful in using the distributed sampler.<|||||>remove the outside wrapping to get to the `transformers`'s model and it should work, 

`model.module.generate()`

You can use the helper that deals with arbitrary number of wrappers.
https://github.com/huggingface/transformers/blob/d90a36d192e2981a41122c30a765c63158dd0557/src/transformers/modeling_utils.py#L3027-L3038

<|||||>> As I already replied earlier `generate` works on multiple gpus including `deepspeed`



> remove the outside wrapping to get to the `transformers`'s model and it should work,
> 
> `model.module.generate()`
> 
> You can use the helper that deals with arbitrary number of wrappers.
> 
> https://github.com/huggingface/transformers/blob/d90a36d192e2981a41122c30a765c63158dd0557/src/transformers/modeling_utils.py#L3027-L3038

I tried unwraping the model as "model.module.generate" but I guess there's an issue in the generators it's all special tokens' index as decoding returns empty strings. 

could you please clarify what model.module is exactly doing?<|||||>@HebaGamalElDin I think your problem is not related to this issue. Or we might need more context here, in your case.
<|||||>> could you please clarify what model.module is exactly do?

Wrappers like DDP, Deepspeed, and others hide the original model inside their objects - usually under `model.module`, so you can get from the wrapped object to the original model using `unwrap_model` which will handle multiple wrappers. If it's just one you can just access the original PreTrainedModel HF subclass model with  `model.module`.

The wrappers have no `generate` method, only the PreTrainedModel subclasses have it. That's why if you want to call `generate` you must call it on the HF model and not its wrappers.

If you have issues I suggest you first remove DDP and debug your issue on a single GPU, once working it'd most likely work under DDP..

As @allanj suggested your issue most likely has nothing to do with unwrapping, so anything is possible if you write your own code. Hence I suggest to sort it out on 1-gpu first, then try 1+.<|||||>@stas00 thank you, what I'm wondering is which model version exactly does model.module access or which model version in which GPU exactly? 

**It was working already in 1 GPU, since I switched to ml.p3.16xlarge instance It makes this behavior.**

Here's the full training process:
```

import os
import torch
import pandas as pd
import random
import math
from copy import deepcopy
from tqdm import tqdm
import os
import re
import shutil
import tarfile, zipfile
import pickle,json
import numpy as np
import itertools
from PIL import Image
import PIL.ImageOps
import cv2
from torch.utils.data import DataLoader
from transformers import AdamW, TrOCRProcessor, VisionEncoderDecoderModel, get_scheduler
from Data_pipeline import Context, HCRDataset, OCRDataLoad
from Validation_Metrics import getWordLevelError, getCharacterLevelError


from datasets import load_metric
cer_metric = load_metric("cer")

# SageMaker data parallel: Import the library PyTorch API
#round(random.uniform(),2)
import smdistributed.dataparallel.torch.torch_smddp

# SageMaker data parallel: Import PyTorch's distributed API
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# SageMaker data parallel: Initialize the process group
dist.init_process_group(backend='smddp')

# LOAD MODEL
def load_model() -> VisionEncoderDecoderModel:
    model: VisionEncoderDecoderModel = VisionEncoderDecoderModel.from_pretrained('gagan3012/ArOCRv4')
    return model.cuda()

# SETUP MODEL CONFIGUATIONS
def init_model_for_training(model: VisionEncoderDecoderModel, processor: TrOCRProcessor):
    model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
    model.config.pad_token_id = processor.tokenizer.pad_token_id
    model.config.vocab_size = model.config.decoder.vocab_size
    model.config.bos_token_id = processor.tokenizer.bos_token_id
    model.config.max_length = 162
    model.config.decoder.is_decoder = True
    model.config.decoder.add_cross_attention = True
    torch.cuda.manual_seed_all(42)

    

def compute_cer(processor, pred_ids, label_ids):
    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id
    label_str = processor.batch_decode(label_ids, skip_special_tokens=True)

    cer = cer_metric.compute(predictions=pred_str, references=label_str)
    return cer


# LOAD PRE_PROCESSOR
def load_processor() -> TrOCRProcessor:
    return TrOCRProcessor.from_pretrained('gagan3012/ArOCRv4')

def train(context: Context, train_epochs, learning_rate):
    model = context.model
    optimizer = AdamW(model.parameters(), lr=learning_rate)
    
    num_training_steps = train_epochs * len(context.train_dataloader)
    lr_scheduler = get_scheduler("linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps)
    
    train_loss = 0.0
    min_cer = 1.0
    min_train_loss = 1.0
    for epoch in range(train_epochs):
        model.train() 
        for j, batch in enumerate(context.train_dataloader):
                inputs: torch.Tensor = batch["input"].cuda(non_blocking=True)
                labels: torch.Tensor = batch["label"].cuda(non_blocking=True)
                #print(inputs)
                #print(labels)
                
                outputs = model(pixel_values=inputs, labels=labels)
                #print(outputs)
                loss = outputs.loss
                loss.backward()
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad(set_to_none=True)
                
                train_loss+=loss
        if (loss < min_train_loss) or (min_train_loss==1.0):
            min_train_loss = loss
        print(f"Epoch {epoch}-----Loss---{train_loss/len(context.train_dataloader)}--------- min-cer: {min_train_loss}")
        
        # evaluate
        #model.eval()
        valid_cer = 0.0
        with torch.no_grad():
            for batch in tqdm(context.val_dataloader):
                #print(f"INPUT: {batch['input']} ------ LABEL: {batch['label']}")
                outputs = model.module.generate(batch["input"].cuda(non_blocking=True))
                #print(f"OUTPUTS: {outputs}")
                #print(f"OUTPUTS on CPU: {outputs.cpu().numpy()}")
                cer = compute_cer(context.processor, pred_ids=outputs.detach(), label_ids=batch["label"])
                valid_cer += cer 

        print("Validation CER:", valid_cer / len(context.val_dataloader))

def main():
    batch_size = 64
    train_epochs = 100
    learning_rate = 0.0001
    checkpoints_path = "checkpoints"
    
    # SageMaker data parallel: Scale batch size by world size
    batch_size //= dist.get_world_size()
    batch_size = max(batch_size, 1)

    # Prepare dataset
    #train_dataset = torchvision.datasets.MNIST(...)
    
    processor = load_processor()
    (x_train,y_train),(x_valid,y_valid),(x_test,y_test) = OCRDataLoad()
    train_dataset = HCRDataset(x_train, y_train, processor)

    # SageMaker data parallel: Set num_replicas and rank in DistributedSampler
    train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_dataset,
            num_replicas=dist.get_world_size(),
            rank=dist.get_rank())
    
    train_dataloader = DataLoader(train_dataset, batch_size, shuffle=False, sampler=train_sampler)
    
    val_dataset = HCRDataset(x_valid, y_valid, processor)
    val_sampler = torch.utils.data.distributed.DistributedSampler(
            val_dataset,
            num_replicas=dist.get_world_size(),
            rank=dist.get_rank())
    val_dataloader = DataLoader(val_dataset, batch_size, shuffle=False, sampler=val_sampler)
 
    # SageMaker data parallel: Wrap the PyTorch model with the library's DDP
    model = load_model()
    init_model_for_training(model, processor)
    
    model = DDP(model, find_unused_parameters=True)
    context = Context(model, processor, train_dataset, train_dataloader, val_dataset, val_dataloader)

    
    # SageMaker data parallel: Pin each GPU to a single library process.
    local_rank = os.environ["LOCAL_RANK"] 
    torch.cuda.set_device(int(local_rank))
    model.cuda(int(local_rank))
    train(context, train_epochs, learning_rate)

    # SageMaker data parallel: Save model on master node.
    if dist.get_rank() == 0:
        model.module.save_pretrained(checkpoints_path)

if __name__ == '__main__':
    main()
 
```

@allanj @stas00 <|||||>In the simple case of DDP you basically have multiple GPUs and each gpu just runs its own `generate` on the unwrapped model.

I see that you use SageMaker here, I don't have experience with this environment, perhaps there is something non-standard about it?

I think you want to start testing with a very simple base case reducing your test to just creating DDP and doing `generate` with some hardcoded string - you can remove all the other code including training and dataloaders, etc.

1. create smddp on 2 gpus
2. run `generate` on that model

nothing else. 

perhaps smddp does something different from other wrappers? Once you have a simple repro code then we can tag someone who knows SMDDP better.<|||||>> > 
> 
> @kkavyashankar0009 @JulesGM
> 
> Hi, it seems that model.generatre() does not support DP, but it can be used in DDP. Here is some example code:
> 
> ```
> import torch.multiprocessing as mp
> import torch.distributed as dist
> import argparse
> import torch
> from transformers import get_linear_schedule_with_warmup, BartConfig, BartTokenizer, BartForConditionalGeneration
> 
> parser = argparse.ArgumentParser()
> parser.add_argument('--nodes', type=int, default=1)  # how many nodes (machines) you have
> parser.add_argument('--gpus', type=int, default=-1, help='num gpus per node')
> parser.add_argument('--nr', type=int, default=0, help='ranking within the nodes')
> args = parser.parse_args()
> tokenizer = BartTokenizer.from_pretrained(args.tokenizer_name)
> 
> def test_model_generation(local_gpu_rank, args):
>     set_seed(args.seed)
>     args.rank = args.nr * args.gpus + local_gpu_rank  # compute the rank of the current GPU
>     dist.init_process_group(backend="nccl", init_method="env://", world_size=args.world_size, rank=args.rank)
> 
>     test_data = "../data/" + args.dataset + "/test.txt"
>     print("Processing data: " + test_data, flush=True)
>     config = BartConfig.from_pretrained(args.config_name)
>     bart_ctx = BartForConditionalGeneration.from_pretrained(args.model_name_or_path, config=config)
>     bart_rep = BartForConditionalGeneration.from_pretrained(args.model_name_or_path, config=config)
>     bart_ctx.resize_token_embeddings(len(tokenizer))
>     bart_rep.resize_token_embeddings(len(tokenizer))
> 
>     model = MyModel(bart, config)
>     model_state_dict = torch.load("../output/model/gen.ddp.pt")  
>     model.load_state_dict(model_state_dict, strict=True)  # load model
>     torch.cuda.set_device(local_gpu_rank)  
>     args.device = torch.device("cuda", local_gpu_rank)
>     model.to(args.device)  # move the model to GPU
>     # model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_gpu_rank], find_unused_parameters=True)
>     model.eval()
> 
>     test_dataset = FileDataset(test_data, tokenizer, max_context_len=args.max_ctx_len, max_response_len=args.max_rep_len, dataset=args.dataset)
>     test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset, shuffle=False)
>     test_dataloader = DataLoader(test_dataset, batch_size=args.test_batch_size, num_workers=2, sampler=test_sampler)  # the batch size on each GPU
> 
>     if args.rank == 0:
>         count = 0
>         fw = open("../output/test.responses.txt", "w", encoding="utf-8")  # open file for writting result
>     with torch.no_grad():
>         test_sampler.set_epoch(0)  # keep all data the same on all GPUs, it is usually used in training, I'm not sure if it is necessary in inference
>         for test_data in test_dataloader:
>             for key in test_data.keys():
>                 test_data[key] = test_data[key].to(args.device)
>             outputs = model.bart_model.generate(
>                 input_ids=test_data["input_ids"],
>                 attention_mask=test_data["attention_mask"],
>                 max_length=args.max_rep_len,
>                 no_repeat_ngram_size=3,
>                 num_beams=10,
>             )  # my model contains a BART model in self.bart_model，so I use model.module.bart_model to get it
>             if outputs.size(1) < args.max_rep_len:  # need padding because the lengths from different GPUs may be different
>                 batch_pred_padding = torch.ones((outputs.size(0), args.max_rep_len - outputs.size(1)), dtype=outputs.dtype).cuda()  # use the padding token of BART, and its token id is 1. Be careful with the data type.
>                 outputs = torch.cat([outputs, batch_pred_padding], dim=1)
>             batch_pred = [torch.zeros_like(outputs, dtype=outputs.dtype).cuda() for _ in range(args.world_size)]  # initialized a list for collecting tensors from all GPUs. Be careful with the data type.
>             dist.all_gather(batch_pred, outputs)  # collect data
>             batch_pred = torch.stack(batch_pred, dim=1)  # use stack, take care of the dimension
>             batch_pred = batch_pred.reshape(-1, args.max_rep_len)
>             if args.rank == 0:
>                 batch_out_sentences = tokenizer.batch_decode(batch_pred, skip_special_tokens=True, clean_up_tokenization_spaces=False)  # decode the token id to token
>                 for r in batch_out_sentences:
>                     fw.write(r + "\n")
>                     fw.flush()
>                 count += len(batch_out_sentences)
>                 print(count)
>     if args.rank == 0:
>         fw.close()
>         
> if __name__ == "__main__":
>     if args.gpus < 0:
>         args.gpus = torch.cuda.device_count()
>     args.world_size = args.nodes * args.gpus
>     os.environ['MASTER_ADDR']='localhost'
>     os.environ['MASTER_PORT']='8888'
>     mp.spawn(test_model_generation, nprocs=args.gpus, args=(args, ))
> ```
> 
> Be careful with the data order. It is better to add IDs and write them to the file, which can be used for checking the order.

May have two errors:

1. gather
```python
dist.all_gather(batch_pred, outputs)  # collect data
batch_pred = torch.stack(batch_pred, dim=1)  # use stack, take care of the dimension
batch_pred = batch_pred.reshape(-1, args.max_rep_len)
```
change to
```python
dist.all_gather(batch_pred, outputs)  # collect data
batch_pred = torch.cat(batch_pred, dim=0)
```

2.  DistributedSampler causes more samples
`batch_out_sentences` need to slice:`batch_out_sentences[:total_examples]`<|||||>@gongel Thanks!
For the first problem, please see the following discussion.
For the second problem, I use DistributedEvalSampler provided in https://github.com/SeungjunNah/DeepDeblur-PyTorch/blob/master/src/data/sampler.py to solve the unbalanced data problem. I think your solution is simpler. Thanks!<|||||>> @gongel Thanks! For the first problem, I do not see the difference between torch.stack + reshape and torch.cat. I think the results are the same. For the second problem, I use DistributedEvalSampler provided in https://github.com/SeungjunNah/DeepDeblur-PyTorch/blob/master/src/data/sampler.py to solve the unbalanced data problem. I think your solution is simpler. Thanks!

For  the first problem, you can test the code
```python
import torch
batch_pred = [torch.ones(2,3), torch.zeros(2,3)]
print(batch_pred)
print(torch.stack(batch_pred, dim=1).reshape(-1, 3))
print(torch.cat(batch_pred, dim=0))
```<|||||>@gongel Thanks. I'm sorry for the mistake.
In this case, I think torch.cat is incorrect. 
In PyTorch DDP, if there are four samples [0, 1, 2 ,3] and two GPUs, the first GPU will get [0, 2] and the second will get [1, 3]. So, using torch.stack+reshape, you will get [0, 1, 2, 3]. Using torch.cat, you will get [0, 2, 1, 3], which is incorrect.<|||||>@DaoD thanks, PyTorch DDP is different from PaddlePaddle DDP.