From 8178a06b90c4683d916a2ec948004338f93427ac Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 10 Sep 2025 20:03:14 +0200 Subject: [PATCH] do detailed import --- src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 6f1ed5173..353549c52 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -22,9 +22,11 @@ from typing import Literal import torch import torch.nn.functional as F # noqa: N812 from torch import Tensor, nn -from transformers import AutoTokenizer, GemmaForCausalLM, PaliGemmaForConditionalGeneration +from transformers import AutoTokenizer from transformers.models.auto import CONFIG_MAPPING from transformers.models.gemma import modeling_gemma +from transformers.models.gemma.modeling_gemma import GemmaForCausalLM +from transformers.models.paligemma.modeling_paligemma import PaliGemmaForConditionalGeneration from lerobot.constants import ACTION, OBS_STATE from lerobot.policies.normalize import Normalize, Unnormalize