Update

[ghstack-poisoned]
pytorch · Dec 20, 2024 · fafa7bd · fafa7bd
2 parents 40390f5 + bde9f05
commit fafa7bd
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -4996,8 +4996,14 @@ def __init__(
                 )
             kwargs = primers
         if not isinstance(kwargs, Composite):
-            kwargs = Composite(**kwargs)
-        self.primers = kwargs
+            shape = kwargs.pop("shape", None)
+            device = kwargs.pop("device", None)
+            if "batch_size" in kwargs.keys():
+                extra_kwargs = {"batch_size": kwargs.pop("batch_size")}
+            else:
+                extra_kwargs = {}
+            primers = Composite(kwargs, device=device, shape=shape, **extra_kwargs)
+        self.primers = primers
         self.expand_specs = expand_specs
 
         if random and default_value: