I am working with the OAI MRI dataset for knee osteoarthritis classification. Each one of 435 MRIs I got has to be classified to a grade. For each MRI in a folder, there are 160 2D images. I created this function to read the dataset:
def dicom2array(path):
  dicom = pydicom.read_file(path)
  data = dicom.pixel_array
  data = (data - np.min(data)) / (np.max(data) - np.min(data))
  data = cv2.resize(data, (256,256))
  return data

def load_3d_dicom_images(scan_id): #returns an object with shape (160,256,256)
  files = sorted(glob.glob(f"/content/drive/MyDrive/subjects/{scan_id}/*/*"))
  img = np.array([dicom2array(a) for a in files])
  return img

class MyCustomDataset(Dataset):
  def __init__(self, path = "/content/drive/MyDrive/subjects"):
    dataframe = pd.read_csv("/content/drive/MyDrive/KL_grade.csv", sep = ';')
    self.labels = {}
    id = list(dataframe["id"])
    grades = list(dataframe["grade"])
    for i,g in zip(id, grades):
      self.labels[str(i).zfill(5)] = g
    
    self.ids = [a.split("/")[-1] for a in sorted(glob.glob(f"/content/drive/MyDrive/subjects/" + "/*"))]

  def __len__(self):
    return len(self.ids)

  def __getitem__(self, idx):
    imgs = load_3d_dicom_images(self.ids[idx])  
    label = self.labels[self.ids[idx]]

    return torch.tensor(imgs, dtype = torch.float32), torch.tensor(label, dtype = torch.long)

I tried to train the model with a resnet18 from the monai library. When I set the batch_size=10 or more cuda runs out of memory (I run it on google colab). It works with batch_size=1 but when I set it equal to 2 I get this error:
stack expects each tensor to be equal size, but got [163, 256, 256] at entry 0 and [160, 256, 256] at entry 1
Can anybody help me ?