hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

Merge remote-tracking branch 'upstream/main' into fix_pusht_diffusion

Browse Source
This commit is contained in:
Alexander Soare
2024-03-21 10:20:52 +00:00
parent acf1174447 1bd50122be
commit 72d3c3120b
15 changed files with 276 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
lerobot/common/policies/diffusion/policy.py
View File

@@ -9,6 +9,7 @@ from lerobot.common.policies.abstract import AbstractPolicy
from lerobot.common.policies.diffusion.diffusion_unet_image_policy import DiffusionUnetImagePolicy
from lerobot.common.policies.diffusion.model.lr_scheduler import get_scheduler
from lerobot.common.policies.diffusion.model.multi_image_obs_encoder import MultiImageObsEncoder, RgbEncoder
from lerobot.common.utils import get_safe_torch_device
class DiffusionPolicy(AbstractPolicy):
@@ -66,9 +67,8 @@ class DiffusionPolicy(AbstractPolicy):
**kwargs,
)
self.device = torch.device(cfg_device)
if torch.cuda.is_available() and cfg_device == "cuda":
self.diffusion.cuda()
self.device = get_safe_torch_device(cfg_device)
self.diffusion.to(self.device)
self.ema_diffusion = None
self.ema = None