Maybe your model is itself too big. I would run the model with batch size of 1 and see how much memory the model itself needs. That gives us an idea about how much data we can upload without OOM error