tangger/lerobot_piper
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Solve conflicts + pre-commit run -a

Browse Source
This commit is contained in:
Cadene
2024-02-29 23:31:32 +00:00
parent 0b9027f05e
commit ae050d2e94
8 changed files with 26 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
lerobot/common/policies/diffusion.py
View File

@@ -5,7 +5,6 @@ import hydra
import torch
import torch.nn as nn
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
from diffusion_policy.model.common.lr_scheduler import get_scheduler
from diffusion_policy.model.vision.model_getter import get_resnet
from diffusion_policy.model.vision.multi_image_obs_encoder import MultiImageObsEncoder
@@ -128,12 +127,8 @@ class DiffusionPolicy(nn.Module):
out = {
"obs": {
"image": batch["observation", "image"].to(
self.device, non_blocking=True
),
"agent_pos": batch["observation", "state"].to(
self.device, non_blocking=True
),
"image": batch["observation", "image"].to(self.device, non_blocking=True),
"agent_pos": batch["observation", "state"].to(self.device, non_blocking=True),
},
"action": batch["action"].to(self.device, non_blocking=True),
}