Polish random generator.

opendilab · Jul 19, 2024 · f6cbd0d · f6cbd0d
1 parent d1791bf
commit f6cbd0d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/grl/generative_models/random_generator.py b/grl/generative_models/random_generator.py
@@ -211,7 +211,10 @@ def generate_data_from_dict(
             data_size, device, batch_size
         )
     else:
-        raise ValueError(f"Invalid data size: {data_size}")
+        try:
+            return list(data_size)
+        except:
+            raise ValueError(f"Invalid data size: {data_size}")
 
 
 if __name__ == "__main__":