Updated ckpt loading in adversarial_attack

PierreMarza · PierreMarza · commit 46bb3e229581 · 2025-06-26T23:36:17.000+02:00
diff --git a/src/thunder/tasks/adversarial_attack.py b/src/thunder/tasks/adversarial_attack.py
@@ -91,7 +91,7 @@ def adversarial_attack(
             f"Checkpoint file not found at {linear_ckpt_path}. Please train a linear probe first."
         )
 
-    ckpt = torch.load(linear_ckpt_path)
+    ckpt = torch.load(linear_ckpt_path, weights_only=True)
     state_dict = ckpt["task_specific_model"]
     out_features, in_features = state_dict["linear.weight"].shape
 

Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,7 @@ def adversarial_attack(`
`91`	`91`	`f"Checkpoint file not found at {linear_ckpt_path}. Please train a linear probe first."`
`92`	`92`	`)`
`93`	`93`
`94`		`- ckpt = torch.load(linear_ckpt_path)`
	`94`	`+ ckpt = torch.load(linear_ckpt_path, weights_only=True)`
`95`	`95`	`state_dict = ckpt["task_specific_model"]`
`96`	`96`	`out_features, in_features = state_dict["linear.weight"].shape`
`97`	`97`