Skip to content
Snippets Groups Projects
Commit 2c86b2b3 authored by Paul Best's avatar Paul Best
Browse files

fix weights relative path

parent fe2dba1f
No related branches found
No related tags found
No related merge requests found
Show whitespace changes
Inline Side-by-side
Showing
with 1 addition and 1 deletion
run_CNN.py
+ 1
− 1
View file @ 2c86b2b3
...@@ -58,7 +58,7 @@ class Dataset(torch.utils.data.Dataset): ...@@ -58,7 +58,7 @@ class Dataset(torch.utils.data.Dataset):
# prepare model # prepare model
model = models.get[args.specie]['archi'] model = models.get[args.specie]['archi']
model.load_state_dict(torch.load(f"{os.path.dirname(__file__)}/weights/{models.get[args.specie]['weights']}")) model.load_state_dict(torch.load(f".{os.path.dirname(__file__)}/weights/{models.get[args.specie]['weights']}"))
model.eval() model.eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device) model.to(device)
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment