Uploading model.pt

config.json

@@ -9,6 +9,7 @@
     "num_experts": 2,
     "hidden_dim": 256,
     "auto_map": {
-        "AutoModel": "model.py"
+        "AutoModel": "model.EmbeddingMoE",
+        "AutoConfig": "model.EmbeddingMoEConfig"
     }
   }

custom_modules.py

@@ -0,0 +1,9 @@
+from transformers import AutoConfig, AutoModel
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+from .model import EmbeddingMoEConfig, EmbeddingMoE
+
+# Register your model
+CONFIG_MAPPING.register("embedding_moe", EmbeddingMoEConfig)
+AutoConfig.register("embedding_moe", EmbeddingMoEConfig)
+AutoModel.register("embedding_moe", EmbeddingMoE)

model.py

@@ -1,5 +1,6 @@
 import torch
 from torch import nn
+import torch.nn.functional as F
 from transformers import PreTrainedModel, PretrainedConfig, AutoModel
 
 class EmbeddingMoEConfig(PretrainedConfig):