chore: import upstream snapshot with attribution

2026-07-13 12:35:17 +08:00
commit 344816a5d8
136 changed files with 25044 additions and 0 deletions
@@ -0,0 +1,18 @@
+from torch import nn
+
+
+class CategorizedModule(nn.Module):
+    @property
+    def category(self):
+        raise NotImplementedError()
+
+    def check_category(self, category):
+        if category is None:
+            raise RuntimeError('Category is not specified in this checkpoint.\n'
+                               'If this is a checkpoint in the old format, please consider '
+                               'migrating it to the new format via the following command:\n'
+                               'python scripts/migrate.py ckpt <INPUT_CKPT> <OUTPUT_CKPT>')
+        elif category != self.category:
+            raise RuntimeError('Category mismatches!\n'
+                               f'This checkpoint is of the category \'{category}\', '
+                               f'but a checkpoint of category \'{self.category}\' is required.')