Hannibal046:

Can I put different parts of the model in different gpu?


Yes, you could use model sharding and push different parts of the model (parameters) to specific devices via to('cuda:id') and use the same operation in the forward to push the activation to the right device.
PyTorch Lightning provides a beta of sharded training 6, which might be interesting for you.
CC @williamFalcon for more information