Polish random generator.

opendilab · Jul 19, 2024 · eaea86c · eaea86c
1 parent f6cbd0d
commit eaea86c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/grl/generative_models/random_generator.py b/grl/generative_models/random_generator.py
@@ -212,7 +212,9 @@ def generate_data_from_dict(
         )
     else:
         try:
-            return list(data_size)
+            return lambda batch_size=None: generate_batch_tensor(
+                    list(data_size), device, batch_size
+                )
         except:
             raise ValueError(f"Invalid data size: {data_size}")