Need help, it is simple neural style transfer script, it works by training but after saving and loading model state, i got this error
RuntimeError: stack expects each tensor to be equal size, but got [1, 64, 256, 256] at entry 0 and [1, 128, 128, 128] at entry 1
at save_image(new_img, ‘output.png’)
i do not understand what is wrong with it, i transformed the image size …
import torch
import torch.nn as nn
import torch.optim as optim
from PIL import Image
import torchvision.transforms as transforms
import torchvision.models as models
from torchvision.utils import save_image
device = torch.device("cuda" if torch.cuda.is_available else "cpu")

image_size  = 256

class VGG(nn.Module):

  def __init__(self):
    super(VGG,self).__init__()
    self.chosen_features = ["0","5","10","19","28"]
    self.model = models.vgg19(pretrained=True).features[:29]

  def forward(self,x):
    features = []
    for layer_num,layer  in enumerate(self.model):
      x = layer(x)
      if str(layer_num) in self.chosen_features:
        features.append(x)
    return features

def load_image(name):
  image = Image.open(name)
  image = loader(image).unsqueeze(0).to(device)
  return image


loader = transforms.Compose([
    transforms.Resize((image_size,image_size)),
    transforms.ToTensor(),
])
model = VGG().to(device).eval()

original_img = load_image("and.png")
style_img = load_image("style.jpg")
generated = original_img.clone().requires_grad_(True)
learning_rate = 0.001
alpha = 1
beta = 0.01
optimizer =  optim.Adam([generated],lr=learning_rate)

for step in range(200):
  optimizer.zero_grad()
  generated_features = model(generated)
  original_img_features = model(original_img)
  style_features = model(style_img)
  style_loss = original_loss = 0
  for gen_feature,orig_feature,style_feature in zip(
      generated_features,original_img_features,style_features
  ):
      batch_size,channel,height,width = gen_feature.shape
      original_loss += torch.mean((gen_feature - orig_feature)**2)
      G = gen_feature.view(channel,height * width) @ gen_feature.view(channel,height * width).T
      A = style_feature.view(channel,height * width) @ style_feature.view(channel,height * width).T
      style_loss += torch.mean((G - A)**2)
  total_loss = alpha * original_loss + beta * style_loss
  total_loss.backward()
  optimizer.step()
  if step % 10 == 0:
    print(total_loss)
    save_image(generated,"gen.png")#works
  torch.save({'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict()},
            "saved_model_for_styletransfer.pth")

  
with torch.no_grad(): 
  checkpoint = torch.load("saved_model_for_styletransfer.pth")
  model = VGG().to(device).eval()
  model.load_state_dict(checkpoint['model_state_dict'])
  new_img = model(load_image("and.png"))
  print(new_img)
  save_image(new_img, 'output.png')#fails