config.json

{
  "architectures": [
    "BranchyCausalModel"
  ],
  "auto_map": {
    "AutoConfig": "BranchyModelConfig.BranchyModelConfig",
    "AutoModelForCausalLM": "BranchyModel.BranchyCausalModel"
  },
  "branch_locations": [
    6,
    12,
    18,
    24
  ],
  "branch_number": 4,
  "confidence_metric": "breaking_ties",
  "copy_lm_head": false,
  "head_thresholds": [
    10.0,
    10.0,
    10.0,
    10.0
  ],
  "head_window_size": 512,
  "model_str": "microsoft/phi-2",
  "model_type": "branchy",
  "penalty_weight": 0.9,
  "torch_dtype": "float32",
  "transformers_version": "4.40.2"
}