Hi Rao,
I guess the problem results from your model is too huge to hold for 1 GPU only.
You need to split model directly on different GPU when writing it from scratch, split training data cannot help and check the below thread.
sorry I don’t have experience to write a multi-GPU training model.
Best,
Xiao