From e82b4c946082ccc58ad96992f4e81d6cef514e11 Mon Sep 17 00:00:00 2001
From: Mahi Shafiullah <3000253+notmahi@users.noreply.github.com>
Date: Mon, 3 Feb 2025 23:30:47 -0500
Subject: [PATCH] Increase model size to match diffusion.

---
 .../common/policies/dit_flow/configuration_dit_flow.py    | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/lerobot/common/policies/dit_flow/configuration_dit_flow.py b/lerobot/common/policies/dit_flow/configuration_dit_flow.py
index 3cff1c96..60093024 100644
--- a/lerobot/common/policies/dit_flow/configuration_dit_flow.py
+++ b/lerobot/common/policies/dit_flow/configuration_dit_flow.py
@@ -122,11 +122,11 @@ class DiTFlowConfig(PreTrainedConfig):
 
     # Diffusion Transformer (DiT) parameters.
     frequency_embedding_dim: int = 256
-    hidden_dim: int = 256
+    hidden_dim: int = 512
     num_blocks: int = 6
-    num_heads: int = 8
-    dropout: float = 0.0
-    dim_feedforward: int = 2048
+    num_heads: int = 16
+    dropout: float = 0.1
+    dim_feedforward: int = 4096
     activation: str = "gelu"
 
     # Noise scheduler.