Update example_inference.py

example_inference.py

@@ -3,15 +3,17 @@ import torch, os
 from PIL import Image
 from briarmbg import BriaRMBG
 from utilities import preprocess_image, postprocess_image
+from huggingface_hub import hf_hub_download
 
 def example_inference():
 
-    model_path = f"{os.path.dirname(os.path.abspath(__file__))}/model.pth"
+    model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
     im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
 
     net = BriaRMBG()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     net.load_state_dict(torch.load(model_path, map_location=device))
+    net.to(device)
     net.eval()    
 
     # prepare input