'shape '[1]' is invalid for input of size 2048

I am trying the extract features from images using resnet 152 model. It was working just fine with resnet18 model but when I try to use resnet152 model, it shows this error "shape '[1]' is invalid for input of size 2048" .I am attaching the whole code and error below .

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
from PIL import Image
# Load the pretrained model
model = models.resnet152(pretrained=True)
# Use the model object to select the desired layer
layer = model._modules.get('avgpool')
# Set model to evaluation mode
model.eval()
scaler = transforms.Resize((224, 224))
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
to_tensor = transforms.ToTensor()

def get_vector(image_name):
  # 1. Load the image with Pillow library
  img = Image.open(image_name)
  # 2. Create a PyTorch Variable with the transformed image
  t_img = Variable(normalize(to_tensor(scaler(img))).unsqueeze(0))
  # 3. Create a vector of zeros that will hold our feature vector
  #    The 'avgpool' layer has an output size of 512
  my_embedding = torch.zeros(512)
  # 4. Define a function that will copy the output of a layer
  def copy_data(m, i, o):
    my_embedding.copy_(o.data.reshape(o.data.size(1)))
  # 5. Attach that function to our selected layer
  h = layer.register_forward_hook(copy_data)
  # 6. Run the model on our transformed image
  model(t_img)
  # 7. Detach our copy function from the layer
  h.remove()
  # 8. Return the feature vector
  return my_embedding

pic_one = '/content/dog1.jpg'
pic_two = '/content/dog3.jpg'
pic_one_vector = get_vector(pic_one)
pic_two_vector = get_vector(pic_two)

Error:

RuntimeError                              Traceback (most recent call last)
<ipython-input-7-ca2d66de2d9c> in <module>()
----> 1 pic_one_vector = get_vector(pic_one)
      2 pic_two_vector = get_vector(pic_two)

5 frames
<ipython-input-5-9a5a4f9a4246> in copy_data(m, i, o)
      9     # 4. Define a function that will copy the output of a layer
      10     def copy_data(m, i, o):
 ---> 11       my_embedding.copy_(o.data.reshape(o.data.size(2)))
      12     # 5. Attach that function to our selected layer
      13     h = layer.register_forward_hook(copy_data)

RuntimeError: shape '[1]' is invalid for input of size 2048

why I am getting this error and how can I solved it..??



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source