sree_harsha:

This is for (1) batch.


Does this mean you are using a batch size of 1?
If so, nn.DataParallel cannot be used, as the batch will be chunked in the batch dimension (dim0) and each chunk will be scattered to the corresponding device.



 sree_harsha:

This is more a case of the model itself being large, would it require explicit ModelParallel sort of optimization?


If the model is too large for a single GPU, you could use model parallel, yes.
However, I’m currently unsure which use cases work.
What’s the largest batch size for a single GPU and what batch sizes are you using for the data parallel approach?