Update torchvision_funcs.py

torchvision_funcs.py

@@ -25,7 +25,7 @@ def deeplabv3_remove_bg(img):
     # with open('./deeplabv3_resnet101.pkl', 'rb') as f:
     #     model = pickle.load(f)
     model = deeplabv3_resnet101(pretrained=True)
-    model.classifier = DeepLabHead(num_classes=1)
+    model.classifier = DeepLabHead(2048, num_classes=1)
     model = model.to(device)
     model.eval()