Update README.md

README.md

@@ -10,8 +10,6 @@ Here is a code to create this tiny model:
 import os
 import torch
 
-torch.set_default_dtype(torch.bfloat16)
-
 from transformers import AutoTokenizer, AutoConfig, Lfm2MoeForCausalLM
 
 # # === Step 1: Define tiny model config ===
@@ -29,6 +27,7 @@ config.num_key_value_heads = 4
 config.hidden_size = 16
 config.num_dense_layers = 1
 config.moe_intermediate_size = 16
+config.intermediate_size = 16
 
 # === Step 2: Create model from config ===
 model = Lfm2MoeForCausalLM(config)