Merge branch 'diffusion-policy' of https://github.com/DorsaRoh/diffusers into diffusion-policy

DorsaRoh · DorsaRoh · commit 9a5cfe722020 · 2024-11-01T11:14:27.000-04:00
diff --git a/examples/reinforcement_learning/diffusion_policy.py b/examples/reinforcement_learning/diffusion_policy.py
@@ -95,7 +95,8 @@ def __init__(self, state_dim=5, device="cpu"):
         )
 
         # load pre-trained weights from HuggingFace
-        checkpoint = torch.load(hf_hub_download("dorsar/diffusion_policy", "push_tblock.pt"), weights_only=True, map_location=device)
+        checkpoint = torch.load(hf_hub_download("dorsar/diffusion_policy", "push_tblock.pt"), map_location=device)
+
         self.model.load_state_dict(checkpoint['model_state_dict'])
         self.obs_encoder.load_state_dict(checkpoint['encoder_state_dict'])
         self.obs_projection.load_state_dict(checkpoint['projection_state_dict'])

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