torch.nn.CosineSimilarity --> IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

I have the following code and I am trying to extract a 2048 dimension vector from the pretrained resnet50 network.

import torch
import torchvision
import torchvision.models as models
from PIL import Image
# Load the pretrained model
model = models.resnet50(pretrained=True)

# Use the model object to select the desired layer
layer = model._modules.get('avgpool')

# Set model to evaluation mode

transforms = torchvision.transforms.Compose([
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

def get_vector(image):
    # Create a PyTorch tensor with the transformed image
    t_img = transforms(image)
    # 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)
    my_embedding = torch.zeros(2048)

    # Define a function that will copy the output of a layer
    def copy_data(m, i, o):
        my_embedding.copy_(o.flatten())                 # <-- flatten

    # Attach that function to our selected layer
    h = layer.register_forward_hook(copy_data)
    # Run the model on our transformed image
    with torch.no_grad():                               # <-- no_grad context
        model(t_img.unsqueeze(0))                       # <-- unsqueeze
    # Detach our copy function from the layer
    # Return the feature vector
    return my_embedding
import torch.nn as nn
cos = nn.CosineSimilarity(dim=1, eps=1e-6)
img1 ='1.jpg')
img2 ='2.jpg')
img3 ='3.jpg')
output1 = cos(pic_vector1, pic_vector2)

I get the following error:

IndexError                                Traceback (most recent call last)
<ipython-input-19-88106798854f> in <module>
----> 1 output1 = cos(pic_vector1, pic_vector2)

~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/ in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/ in forward(self, x1, x2)
     74     def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
---> 75         return F.cosine_similarity(x1, x2, self.dim, self.eps)

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

How can I fix this?

Based on the posted code I assume you want to calculate the cosine similarity between my_embedding and another tensor.
Since my_embedding is a 1-dimensional tensor, using nn.CosineSimilarity(dim=1) won’t work and you could try to use dim=0 or make sure that pic_vector* have at least 2 dimensions.