protonx-models
/

protonx-table-detector

table-recognition

Model card Files Files and versions

hoangth2k4 commited on 10 days ago

Commit

f17bbd4

·

verified ·

1 Parent(s): 87c8029

Update README.md

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -98,7 +98,7 @@ class TableDetector:
         self.model.eval()
     def load_model(self, model_path: str):
-        model = pretrained_models.mobilenet_v2(weights="DEFAULT")
         model.classifier[1]  = nn.Linear(in_features=model.classifier[1].in_features, out_features=2)
         model.load_state_dict(torch.load(model_path, map_location=self.device))
         return model

         self.model.eval()
     def load_model(self, model_path: str):
+        model = pretrained_models.mobilenet_v2(weights=None)
         model.classifier[1]  = nn.Linear(in_features=model.classifier[1].in_features, out_features=2)
         model.load_state_dict(torch.load(model_path, map_location=self.device))
         return model