Include NTK measurement in imports

zincware · May 16, 2024 · 8ca062b · 8ca062b
1 parent 22c6cb6
commit 8ca062b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/papyrus/measurements/__init__.py b/papyrus/measurements/__init__.py
@@ -23,6 +23,7 @@
 
 from papyrus.measurements.base_measurement import BaseMeasurement
 from papyrus.measurements.measurements import (
+    NTK,
     Accuracy,
     Loss,
     NTKEigenvalues,
@@ -41,4 +42,5 @@
     NTKMagnitudeDistribution.__name__,
     Loss.__name__,
     Accuracy.__name__,
+    NTK.__name__,
 ]