Upload folder using huggingface_hub

app.py

@@ -22,7 +22,7 @@ print('DEVICE:', device)
 if device == 'cuda': print('CUDA DEVICE:', torch.cuda.get_device_name())
 
 def load_model_without_module(model, model_path):
-    state_dict = torch.load(model_path, map_location=device, weights_only=True)
+    state_dict = torch.load(model_path, map_location=device, weights_only=False)
 
     new_state_dict = {}
     for k, v in state_dict.items():

u2net/evaluate.py

@@ -11,7 +11,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 print('Device:', device)
 
 def load_model(model, model_path):
-    state_dict = torch.load(model_path, map_location=device, weights_only=True)
+    state_dict = torch.load(model_path, map_location=device, weights_only=False)
     model.load_state_dict(state_dict)
     model.eval()
 

u2net/inference.py

@@ -51,8 +51,7 @@ if __name__ == '__main__':
     # ---
     model = U2Net().to(device)
     model = nn.DataParallel(model)
-    state_dict = torch.load(model_path, map_location=device, weights_only=True)
-    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
+    model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
     model.eval()
 
     mask = run_inference(model, image_path, threshold=None)