Commit 0e970cc

authored

fix: pass weights_only=True to torch.load (#203)

1 parent 02ed72b commit 0e970ccCopy full SHA for 0e970cc

File tree

-1

lines changed

-1

lines changed

Lines changed: 1 addition & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -169,7 +169,7 @@ def model_fn(model_dir):`
`169`	`169`	`device = torch.device("cuda" if torch.cuda.is_available() else "cpu")`
`170`	`170`	`model = torch.nn.DataParallel(Net())`
`171`	`171`	`with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:`
`172`		`- model.load_state_dict(torch.load(f))`
	`172`	`+ model.load_state_dict(torch.load(f, weights_only=True))`
`173`	`173`	`return model.to(device)`
`174`	`174`
`175`	`175`

Comments

(0)