Update model.py

model.py

@@ -56,3 +56,5 @@ def load_model(version='c', device='cpu'):
     state_dict = torch.load(weights_path, map_location=device)
     model.load_state_dict(state_dict, strict=False)
     model.eval()
+
+    return model