The OOM error in tensorflow backend: how the batch_size and the system size influence the GPU memory usage?

[AiPacino9]

I fixed batch_size DP_INFER_BATCH_SIZE= 40,000. Each batch of the model will call 40,000 atoms of data for training. Then the GPU memory usage should be about the same, right? But why do I use a larger system (700 atoms, reviously 400 atoms) for training, OOM happens. Could it be that I misunderstood the memory usage of deepmd-kit?
Besides, for the same data, the backend of torch won't run OOM, but tensorflow will, although tf run much faster. Could it be that tf is trading GPU memory for efficiency?

I want to know if there is any relationship between GPU memory usage, system size(atoms number) and batch_size.

[njzjz]

I fixed batch_size DP_INFER_BATCH_SIZE= 40,000. Each batch of the model will call 40,000 atoms of data for training.

This environment variable is for inference, not for training.

The input JSON file controls the batch size during training.

[AiPacino9]

Thanks for your reply, but I did some tests today:
OOM error happens again when DP_INFER_BATCH_SIZE is not defined during training, although batch_size is set to 1 for all systems in the input json file.
When DP_INFER_BATCH_SIZE is set to 20,000, The training is carried out normally.
When DP_INFER_BATCH_SIZE is set to 40,000, OOM error happens again.

I wonder how the batch_size and the system size(atoms number) influence the GPU memory usage?