code: def get_extensions():
      this_dir = os.path.dirname(os.path.abspath(__file__))
      extensions_dir = os.path.join(
          this_dir, "tensormask", "layers", "csrc"
      )
  
      main_source = os.path.join(extensions_dir, "vision.cpp")
      sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"))
      source_cuda = glob.glob(
          os.path.join(extensions_dir, "**", "*.cu")
      ) + glob.glob(os.path.join(extensions_dir, "*.cu"))
  
      sources = [main_source] + sources
  
      extension = CppExtension
  
      extra_compile_args = {"cxx": []}
      define_macros = []
  
      if (
          torch.cuda.is_available() and CUDA_HOME is not None
      ) or os.getenv("FORCE_CUDA", "0") == "1":
          extension = CUDAExtension
          sources += source_cuda
          define_macros += [("WITH_CUDA", None)]
          extra_compile_args["nvcc"] = [
              "-DCUDA_HAS_FP16=1",
              "-D__CUDA_NO_HALF_OPERATORS__",
              "-D__CUDA_NO_HALF_CONVERSIONS__",
              "-D__CUDA_NO_HALF2_OPERATORS__",
          ]
  
          # It's better if pytorch can do this by default ..
          CC = os.environ.get("CC", None)
          if CC is not None:
              extra_compile_args["nvcc"].append("-ccbin={}".format(CC))
  
      sources = [os.path.join(extensions_dir, s) for s in sources]
  
      include_dirs = [extensions_dir]
  
      ext_modules = [
          extension(
              "tensormask._C",
              sources,
              include_dirs=include_dirs,
              define_macros=define_macros,
              extra_compile_args=extra_compile_args,
          )
      ]
  
      return ext_modules
  
imports: import glob
  import os
  from setuptools import find_packages, setup
  import torch
  from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
label_desc: Use image node_id 1 for calling a global function with example usage: get_extensions() and returns: ext_modules