Merge branch 'main' of https://github.com/moverseai/moai

moverseai · Jun 20, 2024 · 9ddcffe · 9ddcffe
2 parents 715db97 + dca8ee5
commit 9ddcffe
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/moai/engine/run_callback.py b/moai/engine/run_callback.py
@@ -129,7 +129,7 @@ def on_train_batch_end(
         Note: The value ``outputs["loss"]`` here will be the normalized value w.r.t ``accumulate_grad_batches`` of the
             loss returned from ``training_step``.
         """
-        if C._MOAI_LOSSES_ in outputs and 'total' in outputs[C._MOAI_LOSSES_]:
+        if C._MOAI_LOSSES_ in outputs and "total" in outputs[C._MOAI_LOSSES_]:
             if losses := toolz.merge(
                 outputs[f"{C._MOAI_LOSSES_}.weighted"],
                 {"total": outputs[f"{C._MOAI_LOSSES_}.total"]},

diff --git a/moai/monads/human/pose/forward_kinematics.py b/moai/monads/human/pose/forward_kinematics.py
@@ -13,11 +13,12 @@
 class ForwardKinematics(torch.nn.Module):
     def __init__(
         self,
-        parents: list,
-        offsets: list = None,
+        parents: typing.Optional[typing.Sequence[int]] = None,
+        offsets: typing.Optional[typing.Sequence[typing.Sequence[float]]] = None,
     ):  # TODO: add a col/row major param to adjust offset slicing
         super().__init__()
-        self.register_buffer("parents", torch.Tensor(parents).int())
+        if parents is not None:
+            self.register_buffer("parents", torch.Tensor(parents).int())
         if offsets is not None:
             log.warning("offset not set, should be passed in the forward method")
             self.register_buffer("offsets", torch.Tensor(offsets))