sree_harsha:

I was wondering, if considering there are two GPUs with say 16 GB + 16 GB, is there a way to distribute inference across both of them (via ModelParallel) or otherwise, to potentially process them this way?


Yes, this would be possible.Here 34 is a simple example of model sharding.
Basically you can push submodules to specific devices and would have to make sure to push the activation in the forward method to the right device.
Let me know, if you get stuck or need more information.