tag algos (#2011)

mosaicml · Feb 27, 2023 · 8bc5315 · 8bc5315
1 parent b2ec204
commit 8bc5315
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/composer/algorithms/low_precision_groupnorm/low_precision_groupnorm.py b/composer/algorithms/low_precision_groupnorm/low_precision_groupnorm.py
@@ -55,6 +55,13 @@ def __init__(self, apply_at: Event = Event.INIT):
         if self.apply_at not in {Event.INIT, Event.AFTER_LOAD}:
             raise ValueError('LowPrecisionGroupNorm only supports application on Event.INIT and Event.AFTER_LOAD.')
 
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(apply_at={self.apply_at})'
+
+    @staticmethod
+    def required_on_load() -> bool:
+        return True
+
     def match(self, event: Event, state: State) -> bool:
         del state  # unused
         return event == self.apply_at

diff --git a/composer/algorithms/low_precision_layernorm/low_precision_layernorm.py b/composer/algorithms/low_precision_layernorm/low_precision_layernorm.py
@@ -70,6 +70,13 @@ def __init__(self, apply_at: Event = Event.INIT):
         if self.apply_at not in {Event.INIT, Event.AFTER_LOAD}:
             raise ValueError('LowPrecisionLayerNorm only supports application on Event.INIT and Event.AFTER_LOAD.')
 
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(apply_at={self.apply_at})'
+
+    @staticmethod
+    def required_on_load() -> bool:
+        return True
+
     def match(self, event: Event, state: State) -> bool:
         del state  # unused
         return event == self.apply_at