admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 02:06:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

add torch.no_grad decorators

Browse Source
This commit is contained in:
Bryson Jones
2025-12-10 13:48:28 -08:00
parent 3b2a4f548c
commit 3a16a002f8
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py
+2
View File
@@ -760,6 +760,7 @@ class MultiTaskDiTPolicy(PreTrainedPolicy):
self._queues["task"] = deque(maxlen=self.config.n_obs_steps) self._queues["task"] = deque(maxlen=self.config.n_obs_steps)
@torch.no_grad()
def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
"""Predict a chunk of actions given environment observations""" """Predict a chunk of actions given environment observations"""
self.eval() self.eval()
@@ -778,6 +779,7 @@ class MultiTaskDiTPolicy(PreTrainedPolicy):
actions = self._generate_actions(batch) actions = self._generate_actions(batch)
return actions return actions
@torch.no_grad()
def select_action(self, batch: dict[str, Tensor]) -> Tensor: def select_action(self, batch: dict[str, Tensor]) -> Tensor:
"""Select a single action given environment observations""" """Select a single action given environment observations"""
if ACTION in batch: if ACTION in batch: