pytorch 向量相似度计算

import torch
import torch.nn as nn
class CosineSimilarity(nn.Module):
    def forward(self,tensor_1,tensor_2):
        norm_tensor_1=tensor_1.norm(dim=-1, keepdim=True)#将维度缩减为1维
        norm_tensor_2=tensor_2.norm(dim=-1, keepdim=True)
        norm_tensor_1=norm_tensor_1.detach().numpy()
        norm_tensor_2=norm_tensor_2.detach().numpy()
        for i,vec in enumerate(norm_tensor_1):
            if vec[0]==0:
                norm_tensor_1[i][0]=1
        for i,vec in enumerate(norm_tensor_2):
            if vec[0]==0:
                norm_tensor_2[i][0]=1
        norm_tensor_1=torch.tensor(norm_tensor_1)
        norm_tensor_2 = torch.tensor(norm_tensor_2)
        normalized_tensor_1 = tensor_1 / norm_tensor_1
        normalized_tensor_2 = tensor_2 / norm_tensor_2
        res1=normalized_tensor_1*normalized_tensor_2
        res=res1.sum(dim=-1).view(-1,1)
#         res2=1-res
#         res2=res2.sum(dim=-1).view(-1,1) 
#         res=torch.cat((res2,res),1)
        return res
tensor_1=torch.randn((3,200))
tensor_2=torch.randn((3,200))
print(tensor_1)
print(tensor_2)
cos=CosineSimilarity()
c=cos(tensor_1,tensor_2)
print(c)