Update app.py
Browse files
app.py
CHANGED
@@ -24,7 +24,7 @@ custom_model_file = io.BytesIO(custom_model_content)
|
|
24 |
custom_model_state_dict = torch.load(custom_model_file, map_location=torch.device('cpu'))
|
25 |
|
26 |
# Create a new instance of your custom model
|
27 |
-
model = torch.hub.load('ultralytics/yolov5', 'custom',
|
28 |
|
29 |
# Load your custom model into the new instance
|
30 |
model.load_state_dict(custom_model_state_dict)
|
|
|
24 |
custom_model_state_dict = torch.load(custom_model_file, map_location=torch.device('cpu'))
|
25 |
|
26 |
# Create a new instance of your custom model
|
27 |
+
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model=efficientnet_model)
|
28 |
|
29 |
# Load your custom model into the new instance
|
30 |
model.load_state_dict(custom_model_state_dict)
|