code: def opt_amp_callback(
      model: OPTForCausalLM, dtype: torch.dtype
  ) -> None:
      """Casts attention and MLP to low precision only; layernorms stay as f32."""
      for block in model.model.decoder.layers:
          block.self_attn.to(dtype)
          block.fc1.to(dtype)
          block.fc2.to(dtype)
      model.lm_head.to(dtype)
  
file_path: serve/examples/large_models/inferentia2/util/inf2_save_split_checkpoints.py
imports: import argparse
  import os
  import torch
  from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
  from transformers.models.opt import OPTForCausalLM
  from transformers_neuronx.module import save_pretrained_split
label_desc: Use image node_id 2 for calling a global function with example usage: opt_amp_callback(model, dtype) without return types