From ae05279ed0707ed992cc75bdb8ce3725b3dc79d1 Mon Sep 17 00:00:00 2001 From: zilliz Date: Fri, 17 Dec 2021 10:51:50 +0800 Subject: [PATCH] fix dimension problem --- pytorch/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch/model.py b/pytorch/model.py index 6a97391..5a3a050 100644 --- a/pytorch/model.py +++ b/pytorch/model.py @@ -28,7 +28,7 @@ class Model(): super().__init__() model_func = getattr(torchvision.models, model_name) self._model = model_func(pretrained=True) - self._model.fc = torch.nn.Identify() + self._model.fc = torch.nn.Identity() self._model.eval() def __call__(self, img_tensor: torch.Tensor):