Spaces:

DawnC
/

PawMatchAI

Running on Zero

DawnC commited on Apr 7

Commit

403f8f7

verified ·

1 Parent(s): 7f04eb4

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -87,7 +87,7 @@ class ModelManager:
                 'ConvNextV2Base_best_model.pth',
                 map_location=self.device
             )
-            self._breed_model.load_state_dict(checkpoint['base_model'], strict=False)
             self._breed_model.eval()
         return self._breed_model

                 'ConvNextV2Base_best_model.pth',
                 map_location=self.device
             )
+            self._breed_model.load_state_dict(checkpoint['model_state_dict'], strict=False)
             self._breed_model.eval()
         return self._breed_model