Register module for fast instantiation

kaseris · Jan 31, 2024 · 2e69aa1 · 2e69aa1
1 parent 91f5dd3
commit 2e69aa1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/skelcast/losses/euler_angle_loss.py b/src/skelcast/losses/euler_angle_loss.py
@@ -3,8 +3,10 @@
 import torch.nn as nn
 
 from skelcast.data.human36m.quaternion import qeuler
+from skelcast.losses import LOSSES
 
 
+@LOSSES.register_module()
 class EulerAngleLoss(nn.Module):
  def __init__(self, order="xyz", reduction="mean"):
  super(EulerAngleLoss, self).__init__()