diff --git a/reid/models/resnet.py b/reid/models/resnet.py index aa9495d..b316854 100644 --- a/reid/models/resnet.py +++ b/reid/models/resnet.py @@ -40,7 +40,7 @@ def __init__(self, depth, pretrained=True, cut_at_pooling=False, mo.stride = (1,1) #================append conv for FCN==============================# self.num_features = num_features - self.num_classes = 751 #num_classes + self.num_classes = num_classes self.dropout = dropout out_planes = self.base.fc.in_features self.local_conv = nn.Conv2d(out_planes, self.num_features, kernel_size=1,padding=0,bias=False)