From 85f05adcb9517f6e1fa523a4f197f28247eaa81d Mon Sep 17 00:00:00 2001
From: Nikodem Bartnik <nikodem.bartnik@gmail.com>
Date: Mon, 8 Jun 2026 15:51:09 +0200
Subject: [PATCH] fix bug

---
 src/lerobot/policies/pretrained.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py
index cec2f8f25..ab1f1ea03 100644
--- a/src/lerobot/policies/pretrained.py
+++ b/src/lerobot/policies/pretrained.py
@@ -263,7 +263,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             tags=list(set(tags or []).union({"robotics", "lerobot", model_type})),
             model_name=model_type,
             datasets=dataset_repo_id,
-            base_model=base_model_mapping(model_type, None),
+            base_model=base_model_mapping.get(model_type),
         )
 
         template_card = (