code: def _set_up_virtual_devices():
      """Sets up virtual CPU or GPU devices.
  
      This function has to be called before any TF ops has run, and can only be
      called at such moment. The setting is effective for all test classes and
      methods in this file.
      """
      if tf.test.is_gpu_available():
          device_type = "GPU"
          kwargs = {"memory_limit": 1000}  # 1G ram on each virtual GPU
      else:
          device_type = "CPU"
          kwargs = {}
      physical_devices = tf.config.experimental.list_physical_devices(
          device_type
      )
      tf.config.experimental.set_virtual_device_configuration(
          physical_devices[0],
          [
              tf.config.experimental.VirtualDeviceConfiguration(
                  **kwargs
              )
              for _ in range(NUM_REPLICAS)
          ],
      )
  
file_path: neural-structured-learning/neural_structured_learning/keras/adversarial_regularization_multi_device_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import neural_structured_learning.configs
  from neural_structured_learning.keras import adversarial_regularization
  import numpy
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: _set_up_virtual_devices() without return types