Spaces:

ttoosi
/

Generative_Inference

Build error

ttoosi commited on Oct 9, 2024

Commit

b47a80f

1 Parent(s): 741020d

cpu

Files changed (1) hide show

app.py CHANGED Viewed

@@ -10,7 +10,7 @@ checkpoint_path = hf_hub_download(repo_id="ttoosi/resnet50_robust_face", filenam
 # Initialize the model
 model = models.resnet50()
-model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))  # Force model to load on CPU
 model.eval()
 # Image preprocessing

 # Initialize the model
 model = models.resnet50()
+model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu'))['model'])  # Force model to load on CPU
 model.eval()
 # Image preprocessing