OutOfMemory Error on SageMaker with 6 Million Data

[sunilnitdgp]

When attempting to run the code on SageMaker with a dataset of 6 million records and using the instance type ml.m5.4xlarge, I encounter an "OutOfMemory" issue, resulting in the process being killed by SIGKILL (signal 9). This occurs due to high memory consumption during the execution of the code. However, it's working fine in local mac m1 with tensorflow-macos package