I just begin to learn Pytorch and create my first CNN. The dataset contains 3360 RGB images and I converted them to a [3360, 3, 224, 224]. The data and label are in the \'datase