Hy guys, I want to extract the in_feature(2048) of FC layer, passing an image to resnet50.
def get_vector(image):
    
    layer = model._modules.get('avgpool')
    
    my_embedding = torch.zeros(2048) #2048 is the in_features of FC , output of avgpool
   
    def copy_data(m, i, o):
        my_embedding.copy_(o.data)
        
    
    h = layer.register_forward_hook(copy_data)
    
    model(image)#error
  
    h.remove()
    
    # return the vector
    return my_embedding

from PIL import Image
from torch.utils.data import DataLoader
from torchvision import transforms
from os import listdir
from os import path
import time
import h5py
import pandas as pd

since = time.time()
print("---- START FEATURES EXTRACTION ----")
print

column = ["FlickrID", "Features"]

path = "./train_dataset/train_imgs/"

pathCSV = "./train_dataset/features/img_info_TRAIN.csv"



f_id=[]
features_extr=[]

df = pd.DataFrame(columns=column)


tr=transforms.Compose([transforms.Resize(256),
                       transforms.CenterCrop(224),
                       transforms.ToTensor(),
                       transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])


test = Dataset(path, pathCSV, transform=tr)

test_loader = DataLoader(test, batch_size=1, num_workers=2, shuffle = False)
for batch in test_loader:
    nome = batch['FlickrID']
    f_id.append(nome)
    image = batch['image']
    
    
    with torch.no_grad():
        pred = get_vector(image)#error
   features_extr.append(pred)
    
df["FlickrID"] = f_id
df["Features"] = features_extr  


df.to_hdf("Places.h5", key='df', mode='w')

time_elapsed = time.time() - since
print('---- FEATURES SAVED in {:.0f}m {:.0f}s ----'.format(time_elapsed // 60, time_elapsed % 60)) 

I have this error: output with shape [2048] doesn’t match the broadcast shape [1, 2048, 1, 2048]  in the line of code that I have commented with #error
Can you help me?