Update examples/reinforcement_learning/diffusion_policy.py

DorsaRoh · sayakpaul · web-flow · commit 2bad2fc2baf2 · 2024-11-01T11:01:37.000-04:00
Co-authored-by: Sayak Paul &lt;spsayakpaul@gmail.com&gt;
diff --git a/examples/reinforcement_learning/diffusion_policy.py b/examples/reinforcement_learning/diffusion_policy.py
@@ -94,7 +94,7 @@ def __init__(self, state_dim=5, device="cuda" if torch.cuda.is_available() else
         )
 
         # load pre-trained weights from HuggingFace
-        checkpoint = torch.load(hf_hub_download("dorsar/diffusion_policy", "push_tblock.pt"), map_location=device)
+        checkpoint = torch.load(hf_hub_download("dorsar/diffusion_policy", "push_tblock.pt"), weights_only=True, map_location=device)
 
         self.model.load_state_dict(checkpoint['model_state_dict'])
         self.obs_encoder.load_state_dict(checkpoint['encoder_state_dict'])

Original file line number	Diff line number	Diff line change
`@@ -94,7 +94,7 @@ def __init__(self, state_dim=5, device="cuda" if torch.cuda.is_available() else`
`94`	`94`	`)`
`95`	`95`
`96`	`96`	`# load pre-trained weights from HuggingFace`
`97`		`- checkpoint = torch.load(hf_hub_download("dorsar/diffusion_policy", "push_tblock.pt"), map_location=device)`
	`97`	`+ checkpoint = torch.load(hf_hub_download("dorsar/diffusion_policy", "push_tblock.pt"), weights_only=True, map_location=device)`
`98`	`98`
`99`	`99`	`self.model.load_state_dict(checkpoint['model_state_dict'])`
`100`	`100`	`self.obs_encoder.load_state_dict(checkpoint['encoder_state_dict'])`