backup wip

lerobot/common/policies/utils.py

@@ -1,3 +1,7 @@
+import torch
+from torch import nn
+
+
 def populate_queues(queues, batch):
     for key in batch:
         if len(queues[key]) != queues[key].maxlen:
@@ -8,3 +12,21 @@ def populate_queues(queues, batch):
             # add latest observation to the queue
             queues[key].append(batch[key])
     return queues
+
+
+def get_device_from_parameters(module: nn.Module) -> torch.device:
+    """Get a module's device by checking one of its parameters.
+
+    Note: assumes that all parameters have the same device
+    TODO(now): Add test.
+    """
+    return next(iter(module.parameters())).device
+
+
+def get_dtype_from_parameters(module: nn.Module) -> torch.dtype:
+    """Get a module's parameter dtype by checking one of its parameters.
+
+    Note: assumes that all parameters have the same dtype.
+    TODO(now): Add test.
+    """
+    return next(iter(module.parameters())).dtype