Update modeling_i3.py

modeling_i3.py

@@ -28,3 +28,11 @@ class I3ForCausalLM(PreTrainedModel):
     def generate(self, input_ids, max_new_tokens=50, temperature=1.0, top_k=None):
         # Generation method
         return self.model.generate(input_ids, max_new_tokens, temperature, top_k)
+
+from transformers import AutoConfig, AutoModelForCausalLM
+from configuration_i3 import I3Config
+from modeling_i3 import I3ForCausalLM
+
+# Register custom model
+AutoConfig.register("i3", I3Config)
+AutoModelForCausalLM.register(I3Config, I3ForCausalLM)