Not sure how to treat each point within a batch separately

Hey, I am working with PyTorch and cvxpylayers together. I am training a model which first predicts a value from an image (regression) and then uses that value to solve an optimization problem defined in cvxpylayers. The gradients are fed backward after solving the optimization.
I use a batch size of 64, which seemed to be good for the first half (regression from image) but then for some reason the whole 64 size batch of predicted values is inputted to the cvxpylayer, which is only meant to read 1 input.
Any ideas on how I can fix my method to resolve this?

Is it possible for you provide a minimal, run-able code snippet so that we can see how your code is structured?