Computer Vision model

I am trying to implement a computer vision model that I made on Google collabs over into Sagemaker however, the issue that I run into is some type of memory issue possibly caused by the way i upload my training data straight in jyupter lab causing the code to thing it is examining an image with500 channels rather than 500px. The code runs fine on google collabs.