Hey @seungjun, yep, I confirm that your 1st option should work as DDP only calls the default forward() function.

  

      github.com
  

  
    pytorch/pytorch/blob/a33d3d84dfcaef93c17d78cf7983d62ee01d28e3/torch/nn/parallel/distributed.py#L915-L919 1



      
      
          if self.device_ids:
              inputs, kwargs = self.to_kwargs(inputs, kwargs, self.device_ids[0])
              output = self.module(*inputs[0], **kwargs[0])
          else:
              output = self.module(*inputs, **kwargs)
      
    
  


  

  
    
    
  

  


However, there is a caveat. DDP has internal states that requires alternating forward and backward. So if you call things like foward, forward, backward, DDP is likely to hang or crash.