code: def load_pytorch_model(model_name, device="cuda"):
      print("Loading PyTorch model...")
      accelerator = Accelerator()
      tokenizer = AutoTokenizer.from_pretrained(model_name)
      model = AutoModelForCausalLM.from_pretrained(
          model_name,
          torch_dtype=torch.bfloat16,
          trust_remote_code=True,
          device_map="auto",
      )
      model = accelerator.prepare(model)
      model_wrapped = HF_LM(model, tokenizer, device=accelerator.device)
      return model_wrapped
  
file_path: interpret-text/python/interpret_text/generative/model_lib/models.py
imports: import torch
  from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration
  from transformers import LlamaTokenizer, LlamaForCausalLM
  from openai_tooling import CompletionsOpenAI, ChatOpenAI
  from hf_tooling import HF_LM
  from accelerate import Accelerator
label_desc: Use image node_id 1 for calling a global function with example usage: load_pytorch_model(model_name, device) and returns: model_wrapped