config.json

{
  "library_name": "jax",
  "model_type": "custom",
  "architectures": ["SimpleAttentionNetwork"]
}