Adding other files to copy in modeling

modeling_flexbert.py

@@ -116,6 +116,10 @@ from .mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
 from .normalization import get_norm_layer
 from .padding import pad_input, unpad_input
 
+# TODO: This is not used here, but this is so these files are copied when saving the model in ST/PyLate
+from .utils import StrEnum
+from .rotary import UnpaddedRotaryEmbedding
+
 logger = logging.getLogger(__name__)