GPU Memory size change from 16GB to 32GB

Hi NV Team,

Please let me confirm the following question.
Now we tried to do the SSD with the pytorch container(19.12-py3) on DGX Station but it seems that this example did not use the full GPU memory. Would you please teach me the way to change the maximum memory size on SSD code?

Best regards.
Kaka