How to run multi-gpu for a part of model?

Question

I"m trying to run multimodal data, whose sample consists of text tokens shape of (64, 512) and an image shape of (3, 256, 256).

Due to Memory issue, i am trying to run single sample, as my batch, while running text tokens in parallel with 4 GPUS, where I pass (16, 512) to each GPU through encoder (i.e. BERT).

I'm trying to get average embedding of all of them. (i.e. (64, 768) -> (1, 768)). Then with this average embedding, afterwards the model proceeds with one GPU, to compute its relationship with the image.

My question is how can I now run/implement only a part of model for multi-gpus? Thank you in advance!