Sorry for the late reply. The error looks to be from the Dataloader, which in turn is coming from the Dataset due to incorrect dimensions. Try using the compose as below,
compose = transforms.Compose([
    Resize(image_axis_minimum_size, image_axis_minimum_size),
    ToTensor()
     ])