Apply suggestions from code review

huggingface · Oct 11, 2024 · 3db58b6 · 3db58b6
1 parent c8ffcd5
commit 3db58b6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/optimum/utils/input_generators.py b/optimum/utils/input_generators.py
@@ -537,7 +537,7 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
             shape = [self.batch_size, self.sequence_length]
         elif input_name == 'timesteps':
             shape = [self.batch_size, self.sequence_length]
-            return self.random_int_tensor(shape=shape, max_value=max_ep_len, framework=framework, dtype=int_dtype)
+            return self.random_int_tensor(shape=shape, max_value=self.max_ep_len, framework=framework, dtype=int_dtype)
 
         return self.random_float_tensor(shape, min_value=-2., max_value=2., framework=framework, dtype=float_dtype)
 

diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py
@@ -87,8 +87,9 @@ class NormalizedTextConfig(NormalizedConfig):
 
 class NormalizedDecisionTransformerConfig(NormalizedConfig):
     # REFERENCE: https://huggingface.co/docs/transformers/model_doc/decision_transformer
-    STATE_DIM = "state_dim"
     ACT_DIM = "act_dim"
+    STATE_DIM = "state_dim"
+
     MAX_EP_LEN = "max_ep_len"
     HIDDEN_SIZE = "hidden_size"