Update run.py

jan-grzybek-ampere · web-flow · commit c08d358f03b5 · 2025-05-07T03:13:18.000+02:00
diff --git a/recommendation/dlrm/run.py b/recommendation/dlrm/run.py
@@ -84,7 +84,7 @@ def run_single_pass(torch_runner, criteo):
         qr_threshold=200,
         md_threshold=200,
     )
-    dlrm.load_state_dict(torch.load(model_path)["state_dict"])
+    dlrm.load_state_dict(torch.load(model_path, weights_only=False)["state_dict"])
 
     runner = PyTorchRunner(dlrm, example_inputs=dataset.get_inputs(), skip_script=True)
 

Original file line number	Diff line number	Diff line change
`@@ -84,7 +84,7 @@ def run_single_pass(torch_runner, criteo):`
`84`	`84`	`qr_threshold=200,`
`85`	`85`	`md_threshold=200,`
`86`	`86`	`)`
`87`		`- dlrm.load_state_dict(torch.load(model_path)["state_dict"])`
	`87`	`+ dlrm.load_state_dict(torch.load(model_path, weights_only=False)["state_dict"])`
`88`	`88`
`89`	`89`	`runner = PyTorchRunner(dlrm, example_inputs=dataset.get_inputs(), skip_script=True)`
`90`	`90`