code: def prepare_env():
      """
      |1. Check whether arguments required by KMS is set.
      |2. Apply patch to torch.
      """
      if APPID is None or APIKEY is None or EHSM_IP is None:
          print(
              "Please set environment variable APPID, APIKEY, ehsm_ip!"
          )
          exit(1)
      generate_primary_key(EHSM_IP, EHSM_PORT)
      global encrypted_primary_key_path
      encrypted_primary_key_path = "./encrypted_primary_key"
      generate_data_key(
          EHSM_IP, EHSM_PORT, encrypted_primary_key_path, 32
      )
      global encrypted_data_key_path
      encrypted_data_key_path = "./encrypted_data_key"
      patch_encryption()
  
imports: from torch.utils.data import DataLoader
  import torch
  import torch.optim
  from contextlib import nullcontext
  import torch.nn.functional
  import os
  import torch.nn
  import time
  from bigdl.nano.pytorch.patching import patch_encryption
  from bigdl.nano.pytorch.encryption import EncryptedDataset
  import logging
  from bigdl.ppml.kms.ehsm.client import generate_primary_key, generate_data_key, get_data_key_plaintext
  from datasets.load import load_from_disk
  import argparse
  from transformers import BertTokenizer, BertModel, AdamW
label_desc: Use image node_id 1 for calling a global function with example usage: prepare_env() without return types